Kris
c45282f094
- 新增登录状态守卫后台任务 - 新增好友申请自动通过规则引擎 - 新增多分辨率UI配置与模板资源 - 新增消息拉取复合游标支持 - 优化发送队列与UI自动化逻辑 - 新增批量发送日志与错误处理 - 优化Docker镜像构建与ptrace初始化 - 新增联系人名称缓存预热
228 lines
7.7 KiB
Python
228 lines
7.7 KiB
Python
"""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
|