天天看点

写个图像模板匹配算法!开挂?

OpenCV中的模板匹配

OpenCV中的模板匹配是支持基于NCC相似度查找的,但是不是很好用,一个主要的原因是查找最大阈值,只能匹配一个,自己比对阈值,又导致无法正确设定阈值范围,所以问题很多。于是我重新写了纯Python版本的NCC图像模板匹配的代码实现了一个Python版本的,简单易用,支持多尺度,跟多进程并行!

主要思想

主要是基于NCC实现的像素相似度计算,这个OpenCV官方的模板匹配也有这中方式像素相似度计算支持,它的公式描述如下:

写个图像模板匹配算法!开挂?

就是参照这个公式,然后基于OpenCV提供的积分图计算函数,实现了NCC相似度比较计算,值在0~1之间,1表示完全相似,0表示完全不相似。

代码实现

我把整个部分搞成了一个类,调用的方法主要是run_match,就可以直接运行,完成模板匹配。大体的功能跟OpenCV实现的模板匹配功能比较相似,改进的地方就是比较方便的实现多个对象匹配的直接输出Box框。该类完整的代码实现如下:

import cv2 as cv
import numpy as np
import time
import concurrent.futures

class NCCTemplateMatch:
    def __init__(self, ref_imgs, target_imgs, scores, tpl_sums, tpl_sqr_sums, target_sums, target_sqr_sums):
        self.ref_imgs = ref_imgs
        self.target_imgs = target_imgs
        self.scores = scores
        self.tpls_sums = tpl_sums
        self.tpls_sqsums = tpl_sqr_sums
        self.target_sums = target_sums
        self.target_sqsums = target_sqr_sums
        self.nms_boxes = []

    def run_match(self):
        num_ps = min(6, len(self.ref_imgs))
        # print("num_ps: ", num_ps)
        start = time.perf_counter()
        with concurrent.futures.ProcessPoolExecutor(num_ps) as executor:
            matched = executor.map(self.ncc_run, self.ref_imgs, self.target_imgs, self.tpls_sums, self.tpls_sqsums, self.target_sums, self.target_sqsums, self.scores)
            self.nms_boxes = list(matched)
        end = time.perf_counter()
        print(f'Finished in {round(end-start, 2)} seconds')

    def ncc_run(self, tpl_gray, target_gray, tpl_sum, tpl_sqsum, target_sum, target_sqsum, score):
        print("run once~~~~")
        th, tw = tpl_gray.shape
        min_step = max(1, min(th // 16, tw // 16))
        h, w = target_gray.shape
        sr = 1 / (th * tw)
        t_s1 = tpl_sum[th, tw]
        t_s1_2 = t_s1 * t_s1 * sr
        t_s1_1 = t_s1 * sr
        t_s2 = tpl_sqsum[th, tw]
        sum_t = np.sqrt(t_s2 - t_s1_2)
        row = 0
        boxes = []
        confidences = []
        while row < (h - th+1):
            col = 0
            while col < (w - tw+1):
                s1 = self.get_block_sum(target_sum, col, row, col + tw, row + th)
                s2 = self.get_block_sum(target_sqsum, col, row, col + tw, row + th)
                sum1 = t_s1_1 * s1
                ss_sqr = s2 - s1 * s1 * sr
                if ss_sqr < 0:  # fix issue, 精度问题
                    ss_sqr = 0.0
                sum2 = sum_t * np.sqrt(ss_sqr)
                sum3 = np.sum(np.multiply(tpl_gray, target_gray[row:row + th, col:col + tw]))
                if sum2 == 0.0:
                    ncc = 0.0
                else:
                    ncc = (sum3 - sum1) / sum2
                if ncc > score:
                    boxes.append([col, row, tw, th])
                    confidences.append(float(ncc))
                    col += tw//2
                else:
                    col += min_step
            row += min_step

        # NMS Process
        nms_indices = cv.dnn.NMSBoxes(boxes, confidences, 0.5, 0.5)
        det_boxes = []
        print(nms_indices)
        for i in range(len(nms_indices)):
            rect_box = boxes[nms_indices[i]]
            det_boxes.append(rect_box)
        return det_boxes

    def get_block_sum(self, integal_img, x1, y1, x2, y2):
        t1 = integal_img[y1, x1]
        t2 = integal_img[y1, x2]
        t3 = integal_img[y2, x1]
        t4 = integal_img[y2, x2]
        s = t4 - t2 - t3 + t1
        return s      

相关的测试与调用代码如下:

print("test ncc......")
tpl_image = cv.imread("D:/images/llk_tpl.png")
target_image = cv.imread("D:/images/llk.jpg")

tpl_gray = cv.cvtColor(tpl_image, cv.COLOR_BGR2GRAY)
target_gray = cv.cvtColor(target_image, cv.COLOR_BGR2GRAY)
tpl_gray = np.float32(tpl_gray / 255.0)
target_gray = np.float32(target_gray / 255.0)
tpl_sum, tpl_sqsum = cv.integral2(tpl_gray)
t_sum, t_sqsum = cv.integral2(target_gray)
matcher = NCCTemplateMatch([tpl_gray], [target_gray], [0.85],
                           [tpl_sum], [tpl_sqsum], [t_sum], [t_sqsum])
matcher.run_match()
for rect_box in matcher.nms_boxes[0]:
    cv.rectangle(target_image, (rect_box[0], rect_box[1]),
                 (rect_box[0]+rect_box[2], rect_box[1]+rect_box[3]), (0, 0, 255), 2, 8, 0)
cv.imshow("result", target_image)
cv.waitKey(0)
cv.destroyAllWindows()      

模板图像:

写个图像模板匹配算法!开挂?

运行结果如下:

写个图像模板匹配算法!开挂?
写个图像模板匹配算法!开挂?