111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import base64
|
|||
|
|
import json
|
|||
|
|
from dataclasses import dataclass
|
|||
|
|
from typing import Any, Iterable, Optional
|
|||
|
|
|
|||
|
|
import requests
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass(frozen=True)
|
|||
|
|
class OcrBox:
|
|||
|
|
# 4 个顶点坐标（相对截图图片坐标系）
|
|||
|
|
points: tuple[tuple[int, int], tuple[int, int], tuple[int, int], tuple[int, int]]
|
|||
|
|
|
|||
|
|
def center(self) -> tuple[int, int]:
|
|||
|
|
xs = [p[0] for p in self.points]
|
|||
|
|
ys = [p[1] for p in self.points]
|
|||
|
|
return (int(sum(xs) / 4), int(sum(ys) / 4))
|
|||
|
|
|
|||
|
|
|
|||
|
|
@dataclass(frozen=True)
|
|||
|
|
class OcrItem:
|
|||
|
|
text: str
|
|||
|
|
box: OcrBox
|
|||
|
|
|
|||
|
|
|
|||
|
|
class UmiClient:
|
|||
|
|
"""
|
|||
|
|
调用 Umi-OCR 的 HTTP API，并将 data_format=dict 的返回解析为 text+box。
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, url: str = "http://127.0.0.1:1224/api/ocr", timeout_s: float = 15.0) -> None:
|
|||
|
|
self.url = url
|
|||
|
|
self.timeout_s = timeout_s
|
|||
|
|
|
|||
|
|
def check_service(self) -> None:
|
|||
|
|
"""
|
|||
|
|
Umi-OCR 没有稳定的 healthz 文档接口，这里用一次轻量请求做连通性检测。
|
|||
|
|
只要能建立连接并返回 JSON（即使是业务错误），就认为服务已启动。
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
r = requests.post(self.url, json={"base64": "", "options": {"data_format": "dict"}}, timeout=3)
|
|||
|
|
_ = r.text # 触发实际请求
|
|||
|
|
except requests.RequestException as e:
|
|||
|
|
raise RuntimeError(f"无法连接 Umi-OCR 服务：{self.url}。请先在 Umi-OCR 中开启 HTTP 服务。") from e
|
|||
|
|
|
|||
|
|
def ocr_bytes(self, image_bytes: bytes) -> list[OcrItem]:
|
|||
|
|
img64 = base64.b64encode(image_bytes).decode("utf-8")
|
|||
|
|
payload = {"base64": img64, "options": {"data_format": "dict"}}
|
|||
|
|
resp = requests.post(self.url, json=payload, timeout=self.timeout_s)
|
|||
|
|
resp.raise_for_status()
|
|||
|
|
data = resp.json()
|
|||
|
|
return self._parse_umi_dict(data)
|
|||
|
|
|
|||
|
|
def _parse_umi_dict(self, data: dict[str, Any]) -> list[OcrItem]:
|
|||
|
|
# 兼容：当返回不是 dict 时直接报错，方便定位
|
|||
|
|
if not isinstance(data, dict):
|
|||
|
|
raise ValueError(f"Umi-OCR 返回非 JSON 对象：{type(data)}")
|
|||
|
|
|
|||
|
|
items: list[OcrItem] = []
|
|||
|
|
data_list = data.get("data", [])
|
|||
|
|
if not isinstance(data_list, list):
|
|||
|
|
raise ValueError(f"Umi-OCR 返回 data 字段不是 list：{json.dumps(data, ensure_ascii=False)[:500]}")
|
|||
|
|
|
|||
|
|
for it in data_list:
|
|||
|
|
if not isinstance(it, dict):
|
|||
|
|
continue
|
|||
|
|
text = str(it.get("text", "")).strip()
|
|||
|
|
box = it.get("box")
|
|||
|
|
pts = _coerce_box_points(box)
|
|||
|
|
if not text or pts is None:
|
|||
|
|
continue
|
|||
|
|
items.append(OcrItem(text=text, box=OcrBox(points=pts)))
|
|||
|
|
return items
|
|||
|
|
|
|||
|
|
def find_text(
|
|||
|
|
self,
|
|||
|
|
target_name: str,
|
|||
|
|
items: Iterable[OcrItem],
|
|||
|
|
*,
|
|||
|
|
exact: bool = True,
|
|||
|
|
case_sensitive: bool = False,
|
|||
|
|
) -> Optional[tuple[int, int]]:
|
|||
|
|
"""
|
|||
|
|
在给定 OCR items 中查找目标文字，返回其在截图坐标系下的中心点 (x, y)。
|
|||
|
|
"""
|
|||
|
|
t = target_name if case_sensitive else target_name.lower()
|
|||
|
|
for item in items:
|
|||
|
|
s = item.text if case_sensitive else item.text.lower()
|
|||
|
|
ok = (s == t) if exact else (t in s)
|
|||
|
|
if ok:
|
|||
|
|
return item.box.center()
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _coerce_box_points(box: Any) -> Optional[tuple[tuple[int, int], tuple[int, int], tuple[int, int], tuple[int, int]]]:
|
|||
|
|
"""
|
|||
|
|
Umi-OCR 的 box 在 data_format=dict 下通常是 4 个点：
|
|||
|
|
[[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
|
|||
|
|
"""
|
|||
|
|
if not (isinstance(box, list) and len(box) == 4):
|
|||
|
|
return None
|
|||
|
|
pts: list[tuple[int, int]] = []
|
|||
|
|
for p in box:
|
|||
|
|
if not (isinstance(p, (list, tuple)) and len(p) == 2):
|
|||
|
|
return None
|
|||
|
|
pts.append((int(p[0]), int(p[1])))
|
|||
|
|
return (pts[0], pts[1], pts[2], pts[3])
|
|||
|
|
|