基于 OpenCV 的多目标模板匹配，并用聚类算法去除冗余匹配框

模板匹配场景

为了实现多目标匹配，我选用 OpenCV 和 SKlearn 两个库中的模板匹配和聚类算法来实现。

问题描述

多目标匹配时，在同一匹配区域内，会出现多个冗余的框。

import cv2
import numpy as np
from sklearn import cluster
import matplotlib.pyplot as plt

# 读取图片和模板
img_rgb = cv2.imread('../images/aima.png')
img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)#转换成灰度图片
template = cv2.imread('../images/aima-template.png', 0)
h, w = template.shape[:2]

# 进行模板匹配时采用的归一化相关系数
res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)

# 设置置信度阈值
threshold = 0.9

# 取匹配程度大于90%的坐标
loc = np.argwhere(res >= threshold)

# 显示匹配框，并用圆圈标记左上角点
for pt in loc[,::-1]:  # 图片和数组的x对应列，y对应行，所以需要前后交换坐标
    bottom_right = (pt[0]+w, pt[1]+h)
    cv2.rectangle(img_rgb, pt, bottom_right, (0, 0, 255), 2)
    cv2.circle(img_rgb, pt, 5, (255, 255, 255), 2)

# 显示图像
cv2.imshow('img_rgb', img_rgb)
cv2.waitKey(0)

原始图片：

解决方案

单目标检测的思路，只保留置信度最大的匹配框。与此不同的是，多目标匹配需要保留多个区域的匹配候选框。为了解决这个问题引进了聚类算法，在每个区中选择置信度最大的匹配框。
我们使用了近邻传播聚类算法来实现自动计算有多少个匹配的对象。

代码如下：

import cv2
import matplotlib.pyplot as plt
import numpy as np
from sklearn import cluster


def filterRec(res, loc, draw=False):
    """ 对同一对象的多个框按位置聚类后，按置信度选最大的一个进行保留。
    :param res: 是 cv2.matchTemplate 返回值
    :param loc: 是 cv2.np.argwhere(res>threshold) 返回值
    :param draw: 是否进行画图显示
    :return: 返回保留的点的列表 pts
    """
    # 进行聚类分析-近邻传播算法
    model = cluster.AffinityPropagation(damping=0.5, max_iter=500, convergence_iter=30, preference=-50).fit(loc)
    y_pred = model.labels_

    if draw:
        # 画图显示样本数据分类情况
        plt.title('AffinityPropagation', fontsize=16)
        plt.scatter(loc[:, 0], loc[:, 1], s=20, c=y_pred, cmap='brg', label='Samples')
        plt.legend()
        plt.show()

    pts = []
    # 使用循环函数提取每个区域
    for i in set(y_pred):
        argj = loc[y_pred == i]
        argi = argj.T
        # 下面需要注意数组切片操作时，索引需要时 tuple 格式，如果是 numpy.array 格式会报错。
        # 选择置信度最大的点的坐标，注意这时的格式是[行,列]。
        # 输出图片中的坐标时需要转换成 [x, y], 用的方法是 pt[::-1]
        pt = argj[np.argmax(res[tuple(argi)])]
        # 每一区域保留一个选框的左上角坐标
        pts.append(pt[::-1])
    return pts


def main():
    # 读取图片和模板
    img_path = './aima'
    tmp_path = img_path + '-template'
    img_rgb = cv2.imread(img_path + '.png')
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)  # 转换成灰度图片
    template = cv2.imread(tmp_path + '.png', 0)
    h, w = template.shape[:2]

    # 进行模板匹配，方法采用的归一化相关系数
    res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)

    # 设置置信度阈值
    threshold = 0.95
    # 取匹配程度大于90%的坐标
    loc = np.argwhere(res >= threshold)
    print(f"loc length: {len(loc)}")

    # 对匹配框进行过滤
    pts = filterRec(res, loc, False)
    print(f"pts length: {len(pts)}")
    # 显示匹配框，并用圆圈标记左上角点
    for pt in pts:  # 图片和数组的 x 对应列，y 对应行，所以需要前后交换坐标
        bottom_right = (pt[0] + w, pt[1] + h)
        cv2.rectangle(img_rgb, pt, bottom_right, (0, 0, 255), 2)
        cv2.circle(img_rgb, pt, 5, (255, 255, 255), 2)

    # 显示图像
    cv2.imshow('img_rgb', img_rgb)
    cv2.waitKey(0)

if __name__ == "__main__":
    main()

检测后的图片