OpenCV图形模板匹配

OpenCV图形模板匹配

# !/user/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import cv2
import numpy as np
from PyQt5.QtCore import QRect, QPoint

def checkPoint(point, rects):
    for rect in rects:
        if rect.contains(QPoint(point[0], point[1])):
            return True
    return False


def optimizePoint(ls):
    rects = []
    pix = 4
    temp = []
    for pt in ls:
        if checkPoint(pt, rects) is False:
            rect = QRect(QPoint(pt[0] - pix, pt[1] - pix),
                         QPoint(pt[0] + pix, pt[1] + pix))
            rects.append(rect)
            temp.append(QPoint(pt[0], pt[1]))
    return temp


def searchTemplate(src_image, template_image, threshold=0.01):
    CapScreen()
    img_gray = cv2.imread(src_image, 0)
    template = cv2.imread(template_image, 0)
    # cv2.imshow("img_gray", img_gray)
    # cv2.imshow("template", template)
    # cv2.waitKey(0)
    print(threshold)
    res = cv2.matchTemplate(img_gray, template, cv2.TM_SQDIFF_NORMED)
    loc = np.where(res <= threshold)
    ls = list(zip(*loc[::-1]))

    count = len(ls)
    if count <= 0 or count > 100:
        print(count)
        return None

    ls = optimizePoint(ls)
    temp = []
    w, h = template.shape[::-1]
    for point in ls:
        rect = QRect(point.x(), point.y(), w, h)
        temp.append(rect.center())

    return temp


def searchTemplateRects(src_image, template_image, threshold=0.01):
    CapScreen()
    img_gray = cv2.imread(src_image, 0)
    template = cv2.imread(template_image, 0)
    # cv2.imshow("img_gray", img_gray)
    # cv2.imshow("template", template)
    # cv2.waitKey(0)
    print(threshold)
    res = cv2.matchTemplate(img_gray, template, cv2.TM_SQDIFF_NORMED)
    loc = np.where(res <= threshold)
    ls = list(zip(*loc[::-1]))

    count = len(ls)
    if count <= 0 or count > 100:
        print(count)
        return None

    ls = optimizePoint(ls)
    temp = []
    w, h = template.shape[::-1]
    for point in ls:
        rect = QRect(point.x(), point.y(), w, h)
        temp.append(rect)

    return temp