WechatOnCloud/bridge/woc_bridge/ui/backends/opencv.py

"""L1 OpenCVBackend：基于 OpenCV 的模板匹配后端。

提供多尺度模板匹配能力，供 LocatorRegistry/Actions 在 find_element 时
优先用图像定位（精确度高），失败再回退几何兜底。

模板缓存采用 LRU（OrderedDict），上限 max_cache_size=20，避免内存膨胀。
所有 numpy 中间变量显式 del，及时释放 GIL 外的内存。

依赖：opencv-python-headless（容器内已安装）、numpy。
"""

from __future__ import annotations

import asyncio
import logging
import os
import time
from collections import OrderedDict
from typing import Optional

import cv2
import numpy as np

logger = logging.getLogger("woc-bridge")


class OpenCVBackend:
    """OpenCV 模板匹配后端。

    Attributes:
        template_dir: 模板根目录（如 /app/bridge/woc_bridge/ui/profiles/templates/wechat_4.0/light）
        max_cache_size: LRU 缓存上限，默认 20
        _cache: OrderedDict[str, np.ndarray] LRU 缓存（key=模板名，value=灰度图）
    """

    # 默认多尺度（spec 4.3 节）
    _DEFAULT_SCALES = [1.0, 0.75, 1.25, 1.5]

    def __init__(
        self,
        template_dir: str,
        max_cache_size: int = 20,
    ) -> None:
        """初始化 OpenCV 后端。

        Args:
            template_dir: 模板根目录绝对路径
            max_cache_size: LRU 缓存上限
        """
        self.template_dir = template_dir
        self.max_cache_size = max_cache_size
        self._cache: OrderedDict[str, np.ndarray] = OrderedDict()

    def _load_template(self, name: str) -> Optional[np.ndarray]:
        """加载模板为灰度图，LRU 缓存。

        命中：move_to_end（最近使用）。
        未命中：cv2.imread(IMREAD_GRAYSCALE)，超限时 popitem(last=False) 删最旧。
        文件不存在或读取失败返回 None。

        Args:
            name: 模板文件名（如 search_box.png）

        Returns:
            灰度图 np.ndarray 或 None
        """
        if name in self._cache:
            self._cache.move_to_end(name)
            return self._cache[name]
        path = os.path.join(self.template_dir, name)
        if not os.path.exists(path):
            logger.warning("[opencv] template not found: %s", path)
            return None
        tpl = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        if tpl is None:
            logger.warning("[opencv] imread failed: %s", path)
            return None
        # LRU 淘汰
        if len(self._cache) >= self.max_cache_size:
            self._cache.popitem(last=False)
        self._cache[name] = tpl
        return tpl

    async def find_template(
        self,
        screenshot_png: bytes,
        template_name: str,
        threshold: float = 0.80,
        scales: Optional[list[float]] = None,
        roi: Optional[tuple[int, int, int, int]] = None,
    ) -> Optional[tuple[int, int, int, int, float]]:
        """在截图中查找模板，返回 (x, y, w, h, confidence) 或 None。

        - 多尺度 [1.0, 0.75, 1.25, 1.5] 逐个处理
        - 每个 scale 下用 cv2.matchTemplate + minMaxLoc
        - 命中（max_val >= threshold）立即返回，置信度最高的优先
        - ROI 偏移补偿：roi=(x, y, w, h)，返回坐标 = ROI 内坐标 + ROI 偏移
        - 显式 del scaled/res 释放内存

        Args:
            screenshot_png: PNG 字节流（来自 backend.screenshot()）
            template_name: 模板文件名
            threshold: 匹配置信度阈值，默认 0.80
            scales: 自定义尺度列表，None 用 _DEFAULT_SCALES
            roi: (x, y, w, h) 限制匹配区域（窗口绝对坐标），None 全图

        Returns:
            (x, y, w, h, confidence) 绝对坐标，或 None
        """
        # 同步 CPU 密集计算交给线程池
        return await asyncio.to_thread(
            self._find_template_sync,
            screenshot_png,
            template_name,
            threshold,
            scales or self._DEFAULT_SCALES,
            roi,
        )

    def _find_template_sync(
        self,
        screenshot_png: bytes,
        template_name: str,
        threshold: float,
        scales: list[float],
        roi: Optional[tuple[int, int, int, int]],
    ) -> Optional[tuple[int, int, int, int, float]]:
        """同步实现：cv2.matchTemplate 多尺度匹配。

        被 find_template 用 asyncio.to_thread 调用，避免阻塞事件循环。
        """
        # 1. 解码截图
        logger.info("[opencv] decode screenshot: size=%d", len(screenshot_png))
        if not screenshot_png:
            logger.warning("[opencv] screenshot_png is empty")
            return None
        arr = np.frombuffer(screenshot_png, dtype=np.uint8)
        img = cv2.imdecode(arr, cv2.IMREAD_GRAYSCALE)
        if img is None:
            logger.warning("[opencv] decode screenshot failed")
            return None
        del arr

        # 2. 加载模板
        tpl = self._load_template(template_name)
        if tpl is None:
            return None
        th, tw = tpl.shape[:2]

        # 3. ROI 裁剪
        roi_x, roi_y = 0, 0
        if roi is not None:
            rx, ry, rw, rh = roi
            # 边界保护
            rx = max(0, rx)
            ry = max(0, ry)
            rw = max(0, rw)
            rh = max(0, rh)
            img_h, img_w = img.shape[:2]
            x1 = min(rx, img_w)
            y1 = min(ry, img_h)
            x2 = min(rx + rw, img_w)
            y2 = min(ry + rh, img_h)
            if x2 <= x1 or y2 <= y1:
                logger.warning(
                    "[opencv] invalid roi %s in image (%dx%d)",
                    roi, img_w, img_h,
                )
                return None
            img = img[y1:y2, x1:x2]
            roi_x, roi_y = x1, y1

        best: Optional[tuple[int, int, int, int, float]] = None
        for scale in scales:
            # 缩放模板
            new_w = int(tw * scale)
            new_h = int(th * scale)
            if new_w < 8 or new_h < 8:
                continue
            # 模板不能比截图大
            if new_w > img.shape[1] or new_h > img.shape[0]:
                continue
            scaled = cv2.resize(tpl, (new_w, new_h), interpolation=cv2.INTER_AREA)
            try:
                res = cv2.matchTemplate(img, scaled, cv2.TM_CCOEFF_NORMED)
                _, max_val, _, max_loc = cv2.minMaxLoc(res)
            except cv2.error as exc:
                logger.debug(
                    "[opencv] matchTemplate failed scale=%.2f: %s",
                    scale, exc,
                )
                continue
            finally:
                del scaled
            del res

            if max_val >= threshold:
                # max_loc 是 ROI 内坐标，转绝对坐标
                abs_x = max_loc[0] + roi_x
                abs_y = max_loc[1] + roi_y
                # 命中即返回（首个满足阈值的尺度）
                logger.info(
                    "[opencv] match %s scale=%.2f conf=%.3f at (%d,%d)",
                    template_name, scale, max_val, abs_x, abs_y,
                )
                return (abs_x, abs_y, new_w, new_h, float(max_val))

            if best is None or max_val > best[4]:
                best = (
                    max_loc[0] + roi_x,
                    max_loc[1] + roi_y,
                    new_w,
                    new_h,
                    float(max_val),
                )

        if best is not None:
            logger.warning(
                "[opencv] no match for %s (best conf=%.3f < threshold=%.2f)",
                template_name, best[4], threshold,
            )
        else:
            logger.warning(
                "[opencv] no match for %s (no valid scale returned result)",
                template_name,
            )
        return None