PixelSearch.py 文件源码

def find_subimage_in_array(self, sub_image, main_image, threshold=0.40, value=False, debug=False):
        """
        http://docs.opencv.org/3.1.0/d4/dc6/tutorial_py_template_matching.html

        Args:
            sub_image: A numby matrix containing the template we are trying to match
            main_image: A numpy array containing the main image we are trying to find the template in
            value: If true: Similarity is sent back.
            threshold: A treshhold regarding hos sensitive the matching should be.
        Returns:
            A list containing touples:
                If value is true:
                    The touples got he following elements(left,top,right,down,similarity)
                    Where similarity is a measure toward one
                Else:
                    The touples got he following elements(left,top,right,down)

        """
        # TODO: Check the test_init_wnd test for how to implement this :)
        logging.debug("Doing a template match with {} as threshold".format(threshold))
        methods = [cv2.TM_CCOEFF, cv2.TM_CCOEFF_NORMED, cv2.TM_CCORR, cv2.TM_CCORR_NORMED, cv2.TM_SQDIFF,
                   cv2.TM_SQDIFF_NORMED]
        method = methods[0]

        h, w = sub_image.shape[0:2]

        res = cv2.matchTemplate(main_image, sub_image, method)

        loc = np.where(res >= threshold)
        locations = []
        for pt in zip(*loc[::-1]):
            if value:
                locations.append((pt[0], pt[1], pt[0] + w, pt[1] + h, res[pt[1], pt[0]]))
            else:
                locations.append((pt[0], pt[1], pt[0] + w, pt[1] + h))

        logging.debug("Found {} locations".format(len(locations)))
        if debug:
            plt.subplot(121), plt.imshow(res, cmap='gray')
            plt.title('Matching Result'), plt.xticks([]), plt.yticks([])
            plt.subplot(122), plt.imshow(main_image, cmap='gray')
            plt.title('Detected Point'), plt.xticks([]), plt.yticks([])
            for pt in zip(*loc[::-1]):
                cv2.rectangle(main_image, pt, (pt[0] + w, pt[1] + h), (255, 0, 255), 2)
            plt.imshow(main_image)
            plt.show()

        if value:
            locations.sort(reverse=True, key=operator.itemgetter(4))
        return list(map(operator.itemgetter(0, 1, 2, 3), locations))