如何修改OpenCV-Python代码以识别相似图片中的特定形状？

我从img-001.png中截取目标保存为ant0.png，用以下OpenCV-Python代码能在img-001.png中识别多个ant0形状，但在相似图片（如img-002.png）中无法识别，请问该如何修改代码？

原代码：

import cv2
import numpy as np

img_rgb = cv2.imread('img-001.png')

print(img_rgb.shape)
# Convert it to grayscale
img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)
template = cv2.imread('ant0.png',0)
# Store width and heigth of template in w and h
w, h = template.shape[::-1]
# Perform match operations.
res = cv2.matchTemplate(img_gray,template,cv2.TM_CCOEFF_NORMED)
# Specify a threshold
# threshold = 0.84
threshold = 0.8
# Store the coordinates of matched area in a numpy array
loc = np.where(res >= threshold)
x=loc[0]
y=loc[1]
# Draw a rectangle around the matched region.

ant_list=[]
if len(x) and len(y):
    for pt in zip(*loc[::-1]):
        cv2.rectangle(img_rgb, pt, (pt[0] + w, pt[1] + h), (0,255,255), 2)
        print((pt[0] + w, pt[1] + h))
        cv2.imwrite("ant-check0.png", img_rgb)
        ant_list.append(pt[0])
        print('I found the ant')

else:
    print('there is no ant')

print(len(ant_list))
cv2.imshow('ant-check1.png', img_rgb)
cv2.waitKey(0)
cv2.destroyAllWindows()

针对相似图片识别失败的问题，可从以下几个方面修改代码：

1. 更换更鲁棒的模板匹配方法 #

原代码使用cv2.TM_CCOEFF_NORMED，对光照、小形变适应性较弱。可尝试cv2.TM_SQDIFF_NORMED（匹配值越小越相似）或cv2.TM_CCORR_NORMED，并调整阈值逻辑：

# 使用TM_SQDIFF_NORMED时，阈值设为较小值，比如0.2
res = cv2.matchTemplate(img_gray, template, cv2.TM_SQDIFF_NORMED)
threshold = 0.2
loc = np.where(res <= threshold)  # 注意这里是<=判断

2. 加入模板多尺度匹配 #

如果目标在img-002.png中大小有变化，单尺寸模板会失效。可以遍历不同缩放比例，对模板缩放后再匹配：

# 定义缩放范围，比如0.8到1.2，步长0.1
scales = np.arange(0.8, 1.2, 0.1)
for scale in scales:
    # 缩放模板
    resized_template = cv2.resize(template, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    w_scaled, h_scaled = resized_template.shape[::-1]
    # 跳过缩放后尺寸超过原图的情况
    if w_scaled > img_gray.shape[1] or h_scaled > img_gray.shape[0]:
        continue
    res = cv2.matchTemplate(img_gray, resized_template, cv2.TM_CCOEFF_NORMED)
    loc = np.where(res >= threshold)
    for pt in zip(*loc[::-1]):
        cv2.rectangle(img_rgb, pt, (pt[0] + w_scaled, pt[1] + h_scaled), (0,255,255), 2)
        ant_list.append(pt)

3. 优化图像预处理步骤 #

对原图和模板做预处理，减少光照、噪声影响：

# 高斯模糊去噪
img_gray = cv2.GaussianBlur(img_gray, (3,3), 0)
template = cv2.GaussianBlur(template, (3,3), 0)

# 或者直方图均衡化增强对比度
img_gray = cv2.equalizeHist(img_gray)
template = cv2.equalizeHist(template)

4. 调整阈值并加入非极大值抑制（NMS） #

固定阈值容易漏检或误检，可动态调整阈值，同时用NMS去除重复匹配框：

# 动态设置阈值（取匹配最大值的0.7倍）
max_val = np.max(res)
threshold = max_val * 0.7
loc = np.where(res >= threshold)

# 非极大值抑制去重函数
def non_max_suppression(boxes, overlapThresh):
    if len(boxes) == 0:
        return []
    boxes = np.array(boxes)
    x1 = boxes[:,0]
    y1 = boxes[:,1]
    x2 = boxes[:,0] + w
    y2 = boxes[:,1] + h
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(y2)
    pick = []
    while len(idxs) > 0:
        last = len(idxs) - 1
        i = idxs[last]
        pick.append(i)
        xx1 = np.maximum(x1[i], x1[idxs[:last]])
        yy1 = np.maximum(y1[i], y1[idxs[:last]])
        xx2 = np.minimum(x2[i], x2[idxs[:last]])
        yy2 = np.minimum(y2[i], y2[idxs[:last]])
        w_box = np.maximum(0, xx2 - xx1 + 1)
        h_box = np.maximum(0, yy2 - yy1 + 1)
        overlap = (w_box * h_box) / area[idxs[:last]]
        idxs = np.delete(idxs, np.concatenate(([last], np.where(overlap > overlapThresh)[0])))
    return boxes[pick].tolist()

# 收集所有匹配框并去重
boxes = []
for pt in zip(*loc[::-1]):
    boxes.append([pt[0], pt[1]])
filtered_boxes = non_max_suppression(boxes, 0.3)
# 绘制去重后的匹配框
for box in filtered_boxes:
    pt = (box[0], box[1])
    cv2.rectangle(img_rgb, pt, (pt[0] + w, pt[1] + h), (0,255,255), 2)
    ant_list.append(pt[0])

5. 尝试基于特征的匹配（如SIFT） #

如果模板和目标有较大形变，模板匹配效果差，可改用特征点匹配：

# 初始化SIFT检测器
sift = cv2.SIFT_create()
# 检测特征点和描述符
kp1, des1 = sift.detectAndCompute(template, None)
kp2, des2 = sift.detectAndCompute(img_gray, None)
# FLANN匹配器
FLANN_INDEX_KDTREE = 1
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
search_params = dict(checks=50)
flann = cv2.FlannBasedMatcher(index_params, search_params)
matches = flann.knnMatch(des1, des2, k=2)
# 筛选优质匹配点
good = []
for m,n in matches:
    if m.distance < 0.7*n.distance:
        good.append(m)
# 匹配点足够时，计算单应性矩阵并绘制匹配区域
if len(good) > 10:
    src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1,1,2)
    dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1,1,2)
    M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
    h,w = template.shape
    pts = np.float32([[0,0],[0,h-1],[w-1,h-1],[w-1,0]]).reshape(-1,1,2)
    dst = cv2.perspectiveTransform(pts,M)
    img_rgb = cv2.polylines(img_rgb,[np.int32(dst)],True,(0,255,255),2)
else:
    print(f"Not enough matches found - {len(good)}/10")

内容的提问来源于stack exchange，提问作者dominga swopes