80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""屏幕 OCR 识别（RapidOCR，ONNX 推理，跨平台：Windows/Linux/macOS）。
|
||
|
||
用于"条件判断"步骤的 ocr 选择器：截屏 → 识别文字 → 按关键词匹配 → 可点击。
|
||
UI 树里找不到的文字（图片、画布、WebView 渲染的文字）也能识别。
|
||
|
||
依赖 rapidocr_onnxruntime（中英文模型随包内置，无需下载/联网）：
|
||
pip install rapidocr_onnxruntime
|
||
服务器（无显示器）环境建议把 opencv-python 换成 opencv-python-headless。
|
||
|
||
坐标约定：返回 box 为像素坐标（原点左上），与 u2 的 d.click 一致。
|
||
"""
|
||
import threading
|
||
|
||
import numpy as np
|
||
|
||
from core.logger import get_logger
|
||
|
||
_log = get_logger("core.ocr")
|
||
|
||
_engine = None
|
||
_engine_lock = threading.Lock()
|
||
|
||
|
||
def _get_engine():
|
||
global _engine
|
||
with _engine_lock:
|
||
if _engine is None:
|
||
from rapidocr_onnxruntime import RapidOCR
|
||
_log.info("初始化 RapidOCR 引擎（首次约 1-2 秒加载模型）")
|
||
_engine = RapidOCR()
|
||
return _engine
|
||
|
||
|
||
def available():
|
||
"""OCR 是否可用。返回 (ok, hint)。"""
|
||
try:
|
||
_get_engine()
|
||
return True, ""
|
||
except Exception as e:
|
||
return False, f"OCR 不可用（需 pip install rapidocr_onnxruntime）: {e}"
|
||
|
||
|
||
def recognize(image):
|
||
"""识别 PIL Image 上的文字。
|
||
|
||
返回 [{text, score, box}]，box=[x1,y1,x2,y2] 像素坐标（原点左上，与 u2 一致）。
|
||
"""
|
||
engine = _get_engine()
|
||
# PIL(RGB) → ndarray(BGR)，与 OpenCV 约定一致
|
||
arr = np.array(image.convert("RGB"))[:, :, ::-1]
|
||
results = []
|
||
try:
|
||
# 引擎实例并发调用加锁保护（识别本身约 0.2-0.5s，锁开销可忽略）
|
||
with _engine_lock:
|
||
result, _ = engine(arr)
|
||
except Exception as e:
|
||
_log.warning(f"OCR 识别异常: {e}")
|
||
return results
|
||
for item in result or []:
|
||
box4, text, score = item
|
||
xs = [p[0] for p in box4]
|
||
ys = [p[1] for p in box4]
|
||
results.append({"text": text, "score": float(score),
|
||
"box": [round(min(xs)), round(min(ys)),
|
||
round(max(xs)), round(max(ys))]})
|
||
return results
|
||
|
||
|
||
def find_on_screen(image, keyword):
|
||
"""在截图上查找包含 keyword 的文字。
|
||
|
||
返回 (found, center_xy, matched_text)；center_xy 为文字中心像素坐标（可点击），
|
||
未命中返回 (False, None, "")。
|
||
"""
|
||
for r in recognize(image):
|
||
if keyword in r["text"]:
|
||
b = r["box"]
|
||
return True, ((b[0] + b[2]) // 2, (b[1] + b[3]) // 2), r["text"]
|
||
return False, None, ""
|