我只用半小时 | Python手写了个图像模板匹配算法

点击上方蓝字关注我们

微信公众号：OpenCV学堂 关注获取更多计算机视觉与深度学习知识

OpenCV中的模板匹配

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

主要思想

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

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

代码实现

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

代码语言：javascript复制

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

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

代码语言：javascript复制

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()

模板图像：