re-implement equipment with masks

vision.py

@@ -27,32 +27,71 @@ class Vision:
         # TM_CCOEFF, TM_CCOEFF_NORMED, TM_CCORR, TM_CCORR_NORMED, TM_SQDIFF, TM_SQDIFF_NORMED
         self.method = method
 
-    def find(self, haystack_img, needle_img, threshold=0.5, max_results=10, normalize=False, mask=None):
+
+
+    def find_by_mask_and_validate(self, haystack_img, needle_img, needle_mask, max_results=5):
         # run the OpenCV algorithm
         needle_w = needle_img.shape[1]
         needle_h = needle_img.shape[0]
 
-        if normalize:
-            result = cv.matchTemplate(haystack_img, needle_img, cv.TM_CCORR_NORMED, None, mask)
-            _minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
-            cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
-        else:
-            result = cv.matchTemplate(haystack_img, needle_img, self.method)
+        result = cv.matchTemplate(haystack_img, needle_img, cv.TM_CCORR_NORMED, None, needle_mask)
+        cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
+
+        find_num = 20
+        idx_1d = np.argpartition(result.flatten(), -find_num)[-find_num:]
+        idx_2d = np.unravel_index(idx_1d, result.shape)
+
+        rectangles = []
+        for i in range(0, len(idx_2d[0]), 1):
+            y = int(idx_2d[0][i])
+            x = int(idx_2d[1][i])
+            rect = [x, y, needle_w, needle_h]
+            # Add every box to the list twice in order to retain single (non-overlapping) boxes
+            rectangles.append(rect)
+            rectangles.append(rect)
+
+        rectangles, weights = cv.groupRectangles(rectangles, groupThreshold=1, eps=0.5)
+
+        keep_rects = []
+        for rect in rectangles:
+            w = rect[0]
+            h = rect[1]
+            x = rect[2] + w
+            y = rect[3] + h
+            screenshot_pos = haystack_img[h:y, w:x]  # (w, h, x+w, y+h)
+            result2 = cv.matchTemplate(screenshot_pos, needle_img, cv.TM_CCOEFF_NORMED)
+            _minVal2, _maxVal2, minLoc2, maxLoc2 = cv.minMaxLoc(result2, None)
+            #screenshot_pos_img = self.draw_rectangles(screenshot_pos, rectangles)
+            #cv.imshow("screenshot_pos", screenshot_pos)
+            #cv.waitKey(150)
+            if _maxVal2 >= 0.9:
+                keep_rects.append(rect)
+
+        if len(keep_rects) > max_results:
+            keep_rects = keep_rects[:max_results]
+
+        return keep_rects
+
+    def find(self, haystack_img, needle_img, threshold=0.5, max_results=10):
+        # run the OpenCV algorithm
+        needle_w = needle_img.shape[1]
+        needle_h = needle_img.shape[0]
+
+        result = cv.matchTemplate(haystack_img, needle_img, self.method)
 
         # Get the all the positions from the match result that exceed our threshold
         locations = np.where(result >= threshold)
         locations = list(zip(*locations[::-1]))
         # print(locations)
-        _minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
+
+
+        #_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
 
         # if we found no results, return now. this reshape of the empty array allows us to 
         # concatenate together results without causing an error
         if not locations:
             return np.array([], dtype=np.int32).reshape(0, 4)
 
-        if len(locations) > 5000:
-            return np.array([], dtype=np.int32).reshape(0, 4)
-
         # You'll notice a lot of overlapping rectangles get drawn. We can eliminate those redundant
         # locations by using groupRectangles().
         # First we need to create the list of [x, y, w, h] rectangles