AndroidRemoteController/python/scripts/test_match_pic1_screenshot.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
临时测试脚本：对 pic1.png 与 Screenshot-pic1.png 做多组参数匹配，直到找到匹配并在截图上标出。
用法: 在项目根目录执行
  python python/scripts/test_match_pic1_screenshot.py
或
  cd static/process/GenerateNote && python ../../../python/scripts/test_match_pic1_screenshot.py
输出: 控制台打印最佳参数与坐标，并生成 Screenshot-pic1-marked.png 标出匹配框与中心点。
"""
from __future__ import print_function
import os
import sys
import json

# 项目根目录（本脚本在 python/scripts/）
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
PROJECT_ROOT = os.path.normpath(os.path.join(SCRIPT_DIR, '..', '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

# 默认路径：流程目录下的 tmp
PROCESS_DIR = os.path.join(PROJECT_ROOT, 'static', 'process', 'GenerateNote')
TMP_DIR = os.path.join(PROCESS_DIR, 'tmp')
SCREENSHOT_PATH = os.path.join(TMP_DIR, 'Screenshot-pic1.png')
TEMPLATE_PATH = os.path.join(TMP_DIR, 'pic1.png')
MARKED_OUT_PATH = os.path.join(TMP_DIR, 'Screenshot-pic1-marked.png')

try:
    import cv2
    import numpy as np
except ImportError:
    print('请安装: pip install opencv-python numpy')
    sys.exit(1)

# 动态加载 image-match 模块（文件名带横杠）
import importlib.util
_image_match_path = os.path.join(SCRIPT_DIR, 'image-match.py')
_spec = importlib.util.spec_from_file_location('image_match', _image_match_path)
_image_match = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_image_match)
load_image = _image_match.load_image
multi_scale_template_match = _image_match.multi_scale_template_match
match_by_features = _image_match.match_by_features
HAS_ROMA = getattr(_image_match, 'HAS_ROMA', False)
match_by_roma = getattr(_image_match, 'match_by_roma', lambda *a, **k: None)


def multi_scale_with_correlation(screenshot, template, threshold=0.50, scale_min=0.4, scale_max=1.65, steps=80):
    """多尺度模板匹配并返回最佳相关系数。返回 (x, y, w, h, center_x, center_y, best_val) 或 None。"""
    gray_screen = cv2.cvtColor(screenshot, cv2.COLOR_BGR2GRAY)
    gray_tpl = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
    sh, sw = screenshot.shape[:2]
    t_h, t_w = template.shape[:2]
    best = None
    best_val = threshold
    step = max(0.02, (scale_max - scale_min) / float(steps))
    for scale in np.arange(scale_min, scale_max + step * 0.5, step):
        w = max(8, int(round(t_w * scale)))
        h = max(8, int(round(t_h * scale)))
        if h > sh or w > sw:
            continue
        resized = cv2.resize(gray_tpl, (w, h), interpolation=cv2.INTER_AREA)
        result = cv2.matchTemplate(gray_screen, resized, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        if max_val > best_val:
            best_val = max_val
            x, y = int(max_loc[0]), int(max_loc[1])
            center_x = x + w // 2
            center_y = y + h // 2
            best = (x, y, w, h, center_x, center_y, best_val)
    if best is None:
        return None
    return best


def draw_match(screenshot, x, y, w, h, center_x, center_y, out_path):
    """在截图上画匹配框和中心点并保存。"""
    img = screenshot.copy()
    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)
    cv2.circle(img, (center_x, center_y), 8, (0, 0, 255), 2)
    cv2.putText(img, 'center', (center_x + 10, center_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
    cv2.imwrite(out_path, img)
    print('已保存标注图:', out_path)


def main():
    screenshot_path = SCREENSHOT_PATH
    template_path = TEMPLATE_PATH
    if len(sys.argv) >= 3:
        screenshot_path = sys.argv[1]
        template_path = sys.argv[2]
    if not os.path.isfile(screenshot_path):
        print('截图不存在:', screenshot_path)
        sys.exit(1)
    if not os.path.isfile(template_path):
        print('模板不存在:', template_path)
        sys.exit(1)

    screenshot = load_image(screenshot_path)
    template = load_image(template_path)
    if screenshot is None:
        print('无法加载截图')
        sys.exit(1)
    if template is None:
        print('无法加载模板')
        sys.exit(1)

    sh, sw = screenshot.shape[:2]
    t_h, t_w = template.shape[:2]
    print('截图尺寸: {}x{}, 模板尺寸: {}x{}'.format(sw, sh, t_w, t_h))

    # 1) 尝试 RoMa（若可用），仅当中心在上半区时才采用
    if HAS_ROMA:
        roma_result = match_by_roma(screenshot, template, min_matches=4)
        if roma_result is not None:
            x, y, w, h, center_x, center_y = roma_result
            if center_y <= sh * 0.6:
                print('[RoMa] 匹配成功: center=({}, {})'.format(center_x, center_y))
                draw_match(screenshot, x, y, w, h, center_x, center_y, MARKED_OUT_PATH)
                print('推荐 process.json: 保持 [0.15, 2.0], 方形裁剪如 [0.8, "w"]')
                sys.exit(0)

    # 2) 收集多尺度模板匹配结果（参数网格），优先取「中心在画面上半区」的匹配（格子图通常在上方）
    sh, sw = screenshot.shape[:2]
    upper_ratio = 0.6  # 中心 y < sh * upper_ratio 视为上半区（格子区域）
    candidates = []

    # 3) 多尺度模板匹配：参数网格（缩略图通常 scale 约 0.1～0.4）
    scale_min_list = [0.05, 0.10, 0.15]
    scale_max_list = [1.0, 1.5, 2.0]
    threshold_list = [0.40, 0.46, 0.52]
    best_result = None
    best_corr = 0.0
    best_params = None

    for scale_min in scale_min_list:
        for scale_max in scale_max_list:
            if scale_max <= scale_min:
                continue
            for th in threshold_list:
                r = multi_scale_with_correlation(
                    screenshot, template,
                    threshold=th, scale_min=scale_min, scale_max=scale_max, steps=50
                )
                if r is not None:
                    x, y, w, h, cx, cy, corr = r
                    candidates.append((corr, (x, y, w, h, cx, cy), (scale_min, scale_max, th)))

    # 优先选中心在上半区的最高相关匹配，否则选全局最高相关
    best_result = None
    best_corr = 0.0
    best_params = None
    for corr, rect, params in sorted(candidates, key=lambda t: -t[0]):
        x, y, w, h, cx, cy = rect
        if cy <= sh * upper_ratio and corr > best_corr:
            best_corr = corr
            best_result = rect
            best_params = params
    if best_result is None and candidates:
        best_corr, best_result, best_params = max(candidates, key=lambda t: t[0])
        print('(未在上半区找到匹配，使用全局最高相关；请查看标注图确认是否为底部缩略图)')

    if best_result is None:
        print('未找到任何匹配。可尝试: 降低 threshold、扩大 scale_min~scale_max、或检查模板是否在截图内。')
        sys.exit(1)

    x, y, w, h, center_x, center_y = best_result
    scale_min, scale_max, th = best_params
    print('')
    print('[多尺度模板] 匹配成功 (最佳相关系数 {:.3f}, 中心在上半区)'.format(best_corr))
    print('  中心坐标: ({}, {})'.format(center_x, center_y))
    print('  矩形: x={} y={} w={} h={}'.format(x, y, w, h))
    print('  推荐 process.json inVars:')
    print('    inVars: ["tmp/pic{{idx}}.png", [{}, {}], [0.8, "w"]]'.format(scale_min, scale_max))
    print('  并在 image-match.py 中 fallback 阈值可设为 <= {:.2f}'.format(th))
    print('')
    draw_match(screenshot, x, y, w, h, center_x, center_y, MARKED_OUT_PATH)


if __name__ == '__main__':
    main()