re-implement equipment with masks

cv_test_2.py

@@ -0,0 +1,172 @@
+from __future__ import print_function
+import sys
+import cv2 as cv
+import numpy as np
+from hsvfilter import HsvFilter
+
+use_mask = False
+img = None
+templ = None
+mask = None
+image_window = "Source Image"
+result_window = "Result window"
+match_method = 3
+max_Trackbar = 5
+
+
+def main():
+
+
+    global img
+    global templ
+    img = cv.imread("equip/rings/test_screen.jpg", cv.IMREAD_COLOR)
+    templ = cv.imread("equip/rings/ring_1_32.jpg", cv.IMREAD_COLOR)
+
+    hsv = HsvFilter(13, 40, 85, 135, 255, 255, 0, 0, 55, 53)
+    #img = apply_hsv_filter(img, hsv)
+    #templ = apply_hsv_filter(templ, hsv)
+
+
+    global use_mask
+    use_mask = True
+    global mask
+    mask = cv.imread("equip/rings/ring_1_32-mask.png", cv.IMREAD_COLOR)
+
+
+    cv.namedWindow(image_window, cv.WINDOW_AUTOSIZE)
+    cv.namedWindow(result_window, cv.WINDOW_AUTOSIZE)
+
+    trackbar_label = 'Method: \n 0: SQDIFF \n 1: SQDIFF NORMED \n 2: TM CCORR \n 3: TM CCORR NORMED \n 4: TM COEFF \n 5: TM COEFF NORMED'
+    cv.createTrackbar(trackbar_label, image_window, match_method, max_Trackbar, MatchingMethod)
+
+    MatchingMethod(match_method)
+
+    cv.waitKey(0)
+    return 0
+
+def draw_rectangles(haystack_img, rectangles):
+    # these colors are actually BGR
+    line_color = (0, 255, 0)
+    line_type = cv.LINE_4
+    pic = None
+    for (x, y, w, h) in rectangles:
+        # determine the box positions
+        top_left = (x, y)
+        bottom_right = (x + w, y + h)
+        # draw the box
+        cv.rectangle(haystack_img, top_left, bottom_right, line_color, lineType=line_type)
+
+        #pic = haystack_img[y:y + h, x:x + w]
+
+    return haystack_img
+
+def shift_channel(c, amount):
+    if amount > 0:
+        lim = 255 - amount
+        c[c >= lim] = 255
+        c[c < lim] += amount
+    elif amount < 0:
+        amount = -amount
+        lim = amount
+        c[c <= lim] = 0
+        c[c > lim] -= amount
+    return c
+
+def apply_hsv_filter(original_image, hsv_filter):
+    # convert image to HSV
+    hsv = cv.cvtColor(original_image, cv.COLOR_BGR2HSV)
+
+
+    # add/subtract saturation and value
+    h, s, v = cv.split(hsv)
+    s = shift_channel(s, hsv_filter.sAdd)
+    s = shift_channel(s, -hsv_filter.sSub)
+    v = shift_channel(v, hsv_filter.vAdd)
+    v = shift_channel(v, -hsv_filter.vSub)
+    hsv = cv.merge([h, s, v])
+
+    # Set minimum and maximum HSV values to display
+    lower = np.array([hsv_filter.hMin, hsv_filter.sMin, hsv_filter.vMin])
+    upper = np.array([hsv_filter.hMax, hsv_filter.sMax, hsv_filter.vMax])
+    # Apply the thresholds
+    mask = cv.inRange(hsv, lower, upper)
+    result = cv.bitwise_and(hsv, hsv, mask=mask)
+
+    # convert back to BGR for imshow() to display it properly
+    img = cv.cvtColor(result, cv.COLOR_HSV2BGR)
+
+    return img
+
+def MatchingMethod(param):
+    global match_method
+    match_method = param
+
+    img_display = img.copy()
+
+    method_accepts_mask = (cv.TM_SQDIFF == match_method or match_method == cv.TM_CCORR_NORMED)
+    if (use_mask and method_accepts_mask):
+        result = cv.matchTemplate(img, templ, match_method, None, mask)
+    else:
+        result = cv.matchTemplate(img, templ, match_method)
+    #_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
+    cv.normalize(result, result, 0, 1, cv.NORM_MINMAX, -1)
+    #_minVal, _maxVal, minLoc, maxLoc = cv.minMaxLoc(result, None)
+
+    locations = np.where(result >= 0.91)
+    locations = list(zip(*locations[::-1]))
+
+    needle_w = templ.shape[1]
+    needle_h = templ.shape[0]
+
+    find_num = 20
+    idx_1d = np.argpartition(result.flatten(), -find_num)[-find_num:]
+    #new_res = result.flatten()[idx_1d]
+    #new_res.append()
+    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)
+
+
+    for loc in locations:
+        rect = [int(loc[0]), int(loc[1]), 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 = img_display[h:y, w:x]  # (w, h, x+w, y+h)
+        result2 = cv.matchTemplate(screenshot_pos, templ, 5)
+        _minVal2, _maxVal2, minLoc2, maxLoc2 = cv.minMaxLoc(result2, None)
+        if _maxVal2 >= 0.5:
+            keep_rects.append(rect)
+            print("matching error")
+
+
+
+
+
+
+    res1 = draw_rectangles(img_display, keep_rects)
+    res2 = draw_rectangles(img, keep_rects)
+
+    cv.imshow(image_window, res1)
+    cv.imshow(result_window, res2)
+
+
+
+if __name__ == "__main__":
+    main()