AutoAndroidController/py/img-reg.py

"""
图像匹配模块
功能：识别图片2是否在图片1中，如果在则返回图片2在图片1中的坐标
"""

import cv2
import numpy as np
from pathlib import Path
from typing import Tuple, Optional
import os


def read_image_with_unicode_path(image_path):
    """
    读取包含 Unicode 字符（如中文）的图片路径
    OpenCV 的 imread 在 Windows 上可能无法直接处理中文路径
    直接使用 imdecode 方法，避免 imread 的警告
    """
    # 将路径转换为绝对路径并规范化
    abs_path = os.path.abspath(str(image_path))
    
    # 直接使用文件读取 + imdecode，避免 imread 在中文路径上的问题
    # 这样可以避免 cv2.imread 输出的警告信息
    try:
        with open(abs_path, 'rb') as f:
            image_data = f.read()
        # 将字节流解码为图片
        img_array = np.frombuffer(image_data, np.uint8)
        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError(f"cv2.imdecode 无法解码图片: {abs_path}")
        return img
    except FileNotFoundError:
        raise FileNotFoundError(f"图片文件不存在: {abs_path}")
    except Exception as e:
        raise ValueError(f"无法读取图片: {abs_path}, 错误: {e}")


def match_image(
    image1_path: str,
    image2_path: str,
    image1_size: Tuple[int, int] = None
) -> Optional[Tuple[int, int, int, int]]:
    """
    识别图片2是否在图片1中，如果在则返回坐标
    
    Args:
        image1_path: 图片1的本地地址（大图/主图）
        image2_path: 图片2的本地地址（模板图/要查找的图）
        image1_size: 图片1的分辨率尺寸 (width, height)，如果提供则会将图片1缩放到该尺寸
    
    Returns:
        如果找到图片2，返回 (x, y, width, height) 表示图片2在图片1中的位置和尺寸
        如果未找到，返回 None
    """
    # 检查文件是否存在
    img1_path = Path(image1_path)
    img2_path = Path(image2_path)
    
    if not img1_path.exists():
        raise FileNotFoundError(f"图片1不存在: {image1_path}")
    if not img2_path.exists():
        raise FileNotFoundError(f"图片2不存在: {image2_path}")
    
    # 读取图片（使用支持中文路径的方法）
    img1 = read_image_with_unicode_path(img1_path)
    img2 = read_image_with_unicode_path(img2_path)
    
    # 如果提供了图片1的尺寸，则缩放图片1
    if image1_size is not None:
        target_width, target_height = image1_size
        img1 = cv2.resize(img1, (target_width, target_height))
    
    # 转换为灰度图（模板匹配通常在灰度图上进行）
    img1_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    img2_gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    
    # 获取模板尺寸
    template_height, template_width = img2_gray.shape
    
    # 如果模板图比主图大，则无法匹配
    if template_height > img1_gray.shape[0] or template_width > img1_gray.shape[1]:
        return None
    
    # 使用模板匹配方法
    # cv2.TM_CCOEFF_NORMED 返回归一化的相关系数，值越大匹配度越高
    result = cv2.matchTemplate(img1_gray, img2_gray, cv2.TM_CCOEFF_NORMED)
    
    # 获取最佳匹配位置
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
    
    # 设置匹配阈值（可以根据需要调整，0.8 表示80%相似度）
    threshold = 0.8
    
    if max_val >= threshold:
        # 找到匹配位置
        top_left = max_loc
        x, y = top_left
        
        # 返回坐标和尺寸 (x, y, width, height)
        return (x, y, template_width, template_height)
    else:
        # 未找到匹配
        return None


def match_image_multiple(
    image1_path: str,
    image2_path: str,
    image1_size: Tuple[int, int] = None,
    threshold: float = 0.8
) -> list:
    """
    识别图片2在图片1中的所有出现位置（可能有多处匹配）
    
    Args:
        image1_path: 图片1的本地地址
        image2_path: 图片2的本地地址
        image1_size: 图片1的分辨率尺寸 (width, height)
        threshold: 匹配阈值，默认0.8
    
    Returns:
        返回所有匹配位置的列表，每个元素为 (x, y, width, height)
    """
    # 检查文件是否存在
    img1_path = Path(image1_path)
    img2_path = Path(image2_path)
    
    if not img1_path.exists():
        raise FileNotFoundError(f"图片1不存在: {image1_path}")
    if not img2_path.exists():
        raise FileNotFoundError(f"图片2不存在: {image2_path}")
    
    # 读取图片（使用支持中文路径的方法）
    img1 = read_image_with_unicode_path(img1_path)
    img2 = read_image_with_unicode_path(img2_path)
    
    # 如果提供了图片1的尺寸，则缩放图片1
    if image1_size is not None:
        target_width, target_height = image1_size
        img1 = cv2.resize(img1, (target_width, target_height))
    
    # 转换为灰度图
    img1_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    img2_gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    
    # 获取模板尺寸
    template_height, template_width = img2_gray.shape
    
    # 如果模板图比主图大，则无法匹配
    if template_height > img1_gray.shape[0] or template_width > img1_gray.shape[1]:
        return []
    
    # 使用模板匹配
    result = cv2.matchTemplate(img1_gray, img2_gray, cv2.TM_CCOEFF_NORMED)
    
    # 找到所有超过阈值的匹配位置
    locations = np.where(result >= threshold)
    matches = []
    
    # 将匹配位置转换为坐标列表
    for pt in zip(*locations[::-1]):  # Switch x and y coordinates
        x, y = pt
        matches.append((x, y, template_width, template_height))
    
    # 去除重复和相近的匹配（非极大值抑制）
    # 简单实现：如果两个匹配位置太近，只保留置信度更高的
    if len(matches) > 1:
        filtered_matches = []
        for match in matches:
            x, y, w, h = match
            is_duplicate = False
            for existing in filtered_matches:
                ex, ey, _, _ = existing
                # 如果两个匹配中心距离小于模板尺寸的一半，认为是重复
                distance = np.sqrt((x - ex) ** 2 + (y - ey) ** 2)
                if distance < min(w, h) / 2:
                    is_duplicate = True
                    break
            if not is_duplicate:
                filtered_matches.append(match)
        matches = filtered_matches
    
    return matches


if __name__ == "__main__":
    # 测试示例
    import sys
    import os
    
    # 设置 UTF-8 编码，确保能正确处理中文路径
    if sys.platform == 'win32':
        import codecs
        sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
        sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')
    
    if len(sys.argv) < 3:
        print("用法: python img-reg.py <图片1路径> <图片2路径> [图片1宽度] [图片1高度]")
        print("示例: python img-reg.py screenshot.png template.png 1920 1080")
        sys.exit(1)
    
    # 处理路径：将正斜杠转换回反斜杠（Windows），并确保路径正确
    img1_path = sys.argv[1].replace('/', os.sep)
    img2_path = sys.argv[2].replace('/', os.sep)
    
    # 确保路径是绝对路径，并规范化路径
    if not os.path.isabs(img1_path):
        img1_path = os.path.abspath(img1_path)
    if not os.path.isabs(img2_path):
        img2_path = os.path.abspath(img2_path)
    
    # 规范化路径（移除多余的斜杠等）
    img1_path = os.path.normpath(img1_path)
    img2_path = os.path.normpath(img2_path)
    
    image1_size = None
    if len(sys.argv) >= 5:
        try:
            width = int(sys.argv[3])
            height = int(sys.argv[4])
            image1_size = (width, height)
        except ValueError:
            print("警告: 无法解析图片1尺寸，将使用原始尺寸")
    
    # 验证文件是否存在
    if not os.path.exists(img1_path):
        raise FileNotFoundError(f"图片1不存在: {img1_path}")
    if not os.path.exists(img2_path):
        raise FileNotFoundError(f"图片2不存在: {img2_path}")
    
    try:
        result = match_image(img1_path, img2_path, image1_size)
        if result:
            x, y, w, h = result
            print(f"找到匹配！坐标: x={x}, y={y}, 宽度={w}, 高度={h}")
            print(f"JSON格式: {{\"x\": {x}, \"y\": {y}, \"width\": {w}, \"height\": {h}}}")
        else:
            print("未找到匹配")
    except Exception as e:
        print(f"错误: {e}")
        sys.exit(1)