219 lines
11 KiB
Python
219 lines
11 KiB
Python
import os
|
||
import cv2
|
||
import numpy as np
|
||
import imutils
|
||
|
||
|
||
# 工具类
|
||
class AiTools(object):
|
||
|
||
@classmethod
|
||
def find_image_in_image(cls, big_image_path, small_image_path, threshold=0.6, use_multiscale=True, scale_range=(0.7, 1.3), scale_steps=10, visualize=True, enable_subpixel=True, preprocess='histogram'):
|
||
# 读取大图和小图
|
||
large_image = cv2.imread(big_image_path)
|
||
small_image = cv2.imread(small_image_path)
|
||
|
||
if large_image is None or small_image is None:
|
||
print(f"无法加载图像，请检查路径是否正确！\n大图路径: {big_image_path}\n小图路径: {small_image_path}")
|
||
return -1, -1
|
||
|
||
# 打印图像尺寸信息
|
||
print(f"大图尺寸: {large_image.shape[1]}x{large_image.shape[0]}")
|
||
print(f"小图尺寸: {small_image.shape[1]}x{small_image.shape[0]}")
|
||
|
||
# 图像预处理函数
|
||
def preprocess_image(image, method='histogram'):
|
||
if method == 'histogram':
|
||
# 直方图均衡化
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
equalized = cv2.equalizeHist(gray)
|
||
return equalized
|
||
elif method == 'blur':
|
||
# 高斯模糊
|
||
blurred = cv2.GaussianBlur(image, (5, 5), 0)
|
||
gray = cv2.cvtColor(blurred, cv2.COLOR_BGR2GRAY)
|
||
return gray
|
||
elif method == 'edge':
|
||
# 边缘检测
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
edges = cv2.Canny(gray, 100, 200)
|
||
return edges
|
||
elif method == 'sharp':
|
||
# 锐化
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
kernel = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
|
||
sharpened = cv2.filter2D(gray, -1, kernel=kernel)
|
||
return sharpened
|
||
else:
|
||
# 默认转为灰度图
|
||
return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 预处理图像
|
||
large_image_gray = preprocess_image(large_image, preprocess)
|
||
small_image_gray = preprocess_image(small_image, preprocess)
|
||
|
||
best_match_val = -1
|
||
best_match_loc = (-1, -1)
|
||
best_scale = 1.0
|
||
best_method = None
|
||
best_method_enum = None
|
||
|
||
# 定义要尝试的匹配方法
|
||
match_methods = [
|
||
(cv2.TM_CCOEFF, 'TM_CCOEFF'),
|
||
(cv2.TM_CCOEFF_NORMED, 'TM_CCOEFF_NORMED'),
|
||
(cv2.TM_CCORR, 'TM_CCORR'), # 添加这个方法
|
||
(cv2.TM_CCORR_NORMED, 'TM_CCORR_NORMED'),
|
||
(cv2.TM_SQDIFF, 'TM_SQDIFF'),
|
||
(cv2.TM_SQDIFF_NORMED, 'TM_SQDIFF_NORMED')
|
||
]
|
||
|
||
if use_multiscale:
|
||
# 多尺度匹配
|
||
small_height, small_width = small_image_gray.shape
|
||
scale_step = (scale_range[1] - scale_range[0]) / scale_steps
|
||
scales = [scale_range[0] + i * scale_step for i in range(scale_steps + 1)]
|
||
print(f"多尺度匹配范围: {scales}")
|
||
|
||
for scale in scales:
|
||
# 调整小图大小
|
||
new_width = int(small_width * scale)
|
||
new_height = int(small_height * scale)
|
||
if new_width < 10 or new_height < 10: # 防止图像太小
|
||
continue
|
||
# 使用INTER_AREA插值方法提高缩放精度
|
||
resized_small = cv2.resize(small_image_gray, (new_width, new_height), interpolation=cv2.INTER_AREA)
|
||
print(f"尝试尺度: {scale}, 调整后小图尺寸: {new_width}x{new_height}")
|
||
|
||
# 尝试多种匹配方法
|
||
for method in [cv2.TM_CCOEFF_NORMED, cv2.TM_CCORR_NORMED, cv2.TM_SQDIFF_NORMED]:
|
||
method_name = 'TM_CCOEFF_NORMED' if method == cv2.TM_CCOEFF_NORMED else 'TM_CCORR_NORMED' if method == cv2.TM_CCORR_NORMED else 'TM_SQDIFF_NORMED'
|
||
result = cv2.matchTemplate(large_image_gray, resized_small, method)
|
||
if method == cv2.TM_SQDIFF_NORMED:
|
||
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
|
||
current_val = 1 - min_val # 转换为类似其他方法的值
|
||
current_loc = min_loc
|
||
else:
|
||
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
|
||
current_val = max_val
|
||
current_loc = max_loc
|
||
|
||
print(f" 方法: {method_name}, 匹配值: {current_val:.6f}")
|
||
if current_val > best_match_val:
|
||
best_match_val = current_val
|
||
best_match_loc = current_loc
|
||
best_scale = scale
|
||
best_method = method_name
|
||
best_method_enum = method
|
||
else:
|
||
# 单一尺度匹配，但尝试多种方法
|
||
for method in [cv2.TM_CCOEFF_NORMED, cv2.TM_CCORR_NORMED, cv2.TM_SQDIFF_NORMED]:
|
||
method_name = 'TM_CCOEFF_NORMED' if method == cv2.TM_CCOEFF_NORMED else 'TM_CCORR_NORMED' if method == cv2.TM_CCORR_NORMED else 'TM_SQDIFF_NORMED'
|
||
result = cv2.matchTemplate(large_image_gray, small_image_gray, method)
|
||
if method == cv2.TM_SQDIFF_NORMED:
|
||
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
|
||
current_val = 1 - min_val # 转换为类似其他方法的值
|
||
current_loc = min_loc
|
||
else:
|
||
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
|
||
current_val = max_val
|
||
current_loc = max_loc
|
||
|
||
print(f"方法: {method_name}, 匹配值: {current_val:.6f}")
|
||
if current_val > best_match_val:
|
||
best_match_val = current_val
|
||
best_match_loc = current_loc
|
||
best_method = method_name
|
||
best_method_enum = method
|
||
|
||
print(f"最佳匹配值: {best_match_val:.6f}") # 打印最佳匹配值，用于调试
|
||
print(f"最佳匹配方法: {best_method}")
|
||
print(f"使用尺度: {best_scale}") # 打印使用的尺度
|
||
print(f"匹配位置(左上角): {best_match_loc}")
|
||
|
||
# 设置一个阈值，只有当匹配度高于这个阈值时才认为找到
|
||
if best_match_val >= threshold:
|
||
# 计算小图在大图中的中心坐标
|
||
top_left = best_match_loc # 左上角坐标
|
||
# 根据最佳尺度调整小图尺寸
|
||
small_height, small_width = small_image_gray.shape
|
||
adjusted_width = small_width * best_scale
|
||
adjusted_height = small_height * best_scale
|
||
|
||
# 亚像素级精确匹配
|
||
if enable_subpixel and best_method_enum is not None:
|
||
# 调整小图到最佳尺度
|
||
resized_small = cv2.resize(small_image_gray, (int(adjusted_width), int(adjusted_height)), interpolation=cv2.INTER_AREA)
|
||
# 重新执行模板匹配以获取更精确的结果
|
||
result = cv2.matchTemplate(large_image_gray, resized_small, best_method_enum)
|
||
|
||
# 在最佳匹配点周围定义一个小区域进行亚像素精确化
|
||
x, y = top_left
|
||
window_size = 5
|
||
x1, x2 = max(0, x - window_size), min(result.shape[1], x + window_size + 1)
|
||
y1, y2 = max(0, y - window_size), min(result.shape[0], y + window_size + 1)
|
||
|
||
# 提取局部区域
|
||
local_result = result[y1:y2, x1:x2]
|
||
|
||
# 拟合二次曲线找到亚像素级最大值
|
||
if best_method_enum == cv2.TM_SQDIFF_NORMED:
|
||
min_val, _, min_loc, _ = cv2.minMaxLoc(local_result)
|
||
subpixel_offset = (min_loc[0] - window_size, min_loc[1] - window_size)
|
||
else:
|
||
_, max_val, _, max_loc = cv2.minMaxLoc(local_result)
|
||
subpixel_offset = (max_loc[0] - window_size, max_loc[1] - window_size)
|
||
|
||
# 计算亚像素级精确位置
|
||
subpixel_x = x + subpixel_offset[0]
|
||
subpixel_y = y + subpixel_offset[1]
|
||
|
||
print(f"亚像素级精确位置: ({subpixel_x:.2f}, {subpixel_y:.2f})")
|
||
top_left = (subpixel_x, subpixel_y)
|
||
|
||
# 使用浮点运算计算中心坐标
|
||
center_x = top_left[0] + adjusted_width / 2
|
||
center_y = top_left[1] + adjusted_height / 2
|
||
print(f"计算得到的中心坐标: ({center_x:.2f}, {center_y:.2f})")
|
||
|
||
# 可视化匹配结果
|
||
if visualize:
|
||
# 在彩色大图上绘制矩形和中心点
|
||
large_image_copy = large_image.copy()
|
||
# 绘制矩形
|
||
if enable_subpixel:
|
||
top_left_int = (int(round(top_left[0])), int(round(top_left[1])))
|
||
adjusted_width_int = int(round(adjusted_width))
|
||
adjusted_height_int = int(round(adjusted_height))
|
||
bottom_right = (top_left_int[0] + adjusted_width_int, top_left_int[1] + adjusted_height_int)
|
||
center_int = (int(round(center_x)), int(round(center_y)))
|
||
else:
|
||
bottom_right = (top_left[0] + int(adjusted_width), top_left[1] + int(adjusted_height))
|
||
center_int = (int(center_x), int(center_y))
|
||
top_left_int = top_left
|
||
|
||
cv2.rectangle(large_image_copy, top_left_int, bottom_right, (0, 255, 0), 2)
|
||
cv2.circle(large_image_copy, center_int, 5, (0, 0, 255), -1)
|
||
# 显示图像
|
||
cv2.imshow('匹配结果', large_image_copy)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows()
|
||
|
||
return center_x, center_y
|
||
else:
|
||
print(f"匹配值 {best_match_val:.6f} 低于阈值 {threshold}，未找到匹配")
|
||
return -1, -1
|
||
|
||
# 根据名称获取图片地址
|
||
@classmethod
|
||
def pathWithName(cls, name):
|
||
current_file_path = os.path.abspath(__file__)
|
||
# 获取当前文件所在的目录（即script目录）
|
||
current_dir = os.path.dirname(current_file_path)
|
||
# 由于script目录位于项目根目录下一级，因此需要向上一级目录移动两次
|
||
project_root = os.path.abspath(os.path.join(current_dir, '..'))
|
||
# 构建资源文件的完整路径，向上两级目录，然后进入 resources 目录
|
||
resource_path = os.path.abspath(os.path.join(project_root, 'resources', name + ".jpg")).replace('/', '\\\\')
|
||
return resource_path
|
||
|