vision.py

"""OpenCV 模板加载与匹配"""

import glob
import os
from dataclasses import dataclass
from typing import List, Tuple

import cv2
import numpy as np


@dataclass
class Template:
    name: str
    image: np.ndarray


def preprocess(image_bgr: np.ndarray, use_edge: bool = True) -> np.ndarray:
    gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
    if use_edge:
        return cv2.Canny(gray, 100, 200)
    return cv2.GaussianBlur(gray, (3, 3), 0)


def load_templates(template_dir: str, use_edge: bool = True) -> List[Template]:
    pattern = os.path.join(template_dir, "*.png")
    paths = sorted(glob.glob(pattern))
    templates: List[Template] = []
    for path in paths:
        raw = cv2.imread(path)
        if raw is None:
            continue
        processed = preprocess(raw, use_edge=use_edge)
        templates.append(Template(name=os.path.basename(path), image=processed))
    if not templates:
        raise FileNotFoundError(f"模板目录 {template_dir} 中没有 PNG 文件")
    return templates


def match_template(
    frame_processed: np.ndarray,
    templates: List[Template],
    target_name: str,
    scale: float = 1.0,
) -> float:
    """匹配指定名称的模板，返回最高分数。"""
    target_key = target_name.strip().lower()
    fh, fw = frame_processed.shape[:2]

    for tpl in templates:
        if tpl.name.strip().lower() != target_key:
            continue
        tpl_img = tpl.image
        if abs(scale - 1.0) > 0.01:
            new_w = max(1, int(tpl_img.shape[1] * scale))
            new_h = max(1, int(tpl_img.shape[0] * scale))
            tpl_img = cv2.resize(tpl_img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        th, tw = tpl_img.shape[:2]
        if th > fh or tw > fw:
            continue
        result = cv2.matchTemplate(frame_processed, tpl_img, cv2.TM_CCOEFF_NORMED)
        _, max_val, _, _ = cv2.minMaxLoc(result)
        return float(max_val)
    return 0.0