from __future__ import annotations import base64 from pathlib import Path from typing import Any import httpx class OCRClientError(RuntimeError): def __init__(self, code: str, message: str, details: Any = None) -> None: super().__init__(f"{code}: {message}") self.code = code self.message = message self.details = details class OCRClient: def __init__( self, base_url: str = "http://127.0.0.1:5020", timeout: float = 120.0 ) -> None: self.base_url = base_url.rstrip("/") self.timeout = timeout @staticmethod def path_image(path: str | Path) -> dict[str, str]: return {"path": str(Path(path).resolve())} @staticmethod def base64_image(content: bytes) -> dict[str, str]: return {"base64": base64.b64encode(content).decode("ascii")} def invoke(self, operation: str, payload: dict[str, Any] | None = None) -> Any: response = httpx.post( f"{self.base_url}/api/invoke", json={"operation": operation, "payload": payload or {}}, timeout=self.timeout, ) try: body = response.json() except ValueError as exc: raise OCRClientError( "INVALID_RESPONSE", f"OCR 服务返回非 JSON 响应: HTTP {response.status_code}", response.text[:500], ) from exc if not body.get("ok"): error = body.get("error") or {} raise OCRClientError( error.get("code", "UNKNOWN_ERROR"), error.get("message", "OCR 服务调用失败"), error.get("details"), ) response.raise_for_status() return body["result"] def status(self) -> dict[str, Any]: return self.invoke("providers.status") def warmup( self, providers: list[str], languages: dict[str, list[str]] | None = None, ) -> dict[str, Any]: return self.invoke( "providers.warmup", {"providers": providers, "languages": languages or {}}, ) def recognize( self, image: dict[str, str], *, provider: str = "rapidocr", languages: list[str] | None = None, min_confidence: float = 0.0, options: dict[str, Any] | None = None, ) -> dict[str, Any]: return self.invoke( "ocr.recognize", { "image": image, "provider": provider, "languages": languages, "min_confidence": min_confidence, "options": options or {}, }, ) def compare( self, image: dict[str, str], *, providers: list[str], languages: list[str] | None = None, min_confidence: float = 0.0, options: dict[str, dict[str, Any]] | None = None, ) -> dict[str, Any]: return self.invoke( "ocr.compare", { "image": image, "providers": providers, "languages": languages, "min_confidence": min_confidence, "options": options or {}, }, )