AutoAndroidController/py/img-reg.py

"""
鍥惧儚鍖归厤妯″潡
鍔熻兘锛氳瘑鍒浘鐗?鏄惁鍦ㄥ浘鐗?涓紝濡傛灉鍦ㄥ垯杩斿洖鍥剧墖2鍦ㄥ浘鐗?涓殑鍧愭爣
"""

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鏄惁鍦ㄥ浘鐗?涓紝濡傛灉鍦ㄥ垯杩斿洖鍧愭爣
    
    Args:
        image1_path: 鍥剧墖1鐨勬湰鍦板湴鍧€锛堝ぇ鍥?涓诲浘锛?        image2_path: 鍥剧墖2鐨勬湰鍦板湴鍧€锛堟ā鏉垮浘/瑕佹煡鎵剧殑鍥撅級
        image1_size: 鍥剧墖1鐨勫垎杈ㄧ巼灏哄 (width, height)锛屽鏋滄彁渚涘垯浼氬皢鍥剧墖1缂╂斁鍒拌灏哄
    
    Returns:
        濡傛灉鎵惧埌鍥剧墖2锛岃繑鍥?(x, y, width, height) 琛ㄧず鍥剧墖2鍦ㄥ浘鐗?涓殑浣嶇疆鍜屽昂瀵?        濡傛灉鏈壘鍒帮紝杩斿洖 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)
    
    # 濡傛灉鎻愪緵浜嗗浘鐗?鐨勫昂瀵革紝鍒欑缉鏀惧浘鐗?
    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鍦ㄥ浘鐗?涓殑鎵€鏈夊嚭鐜颁綅缃紙鍙兘鏈夊澶勫尮閰嶏級
    
    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)
    
    # 濡傛灉鎻愪緵浜嗗浘鐗?鐨勫昂瀵革紝鍒欑缉鏀惧浘鐗?
    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)