"""OCR 插件核心。 对 PaddleOCR 3.x predict 接口做一层薄封装:解析识别结果、按需写 run 事件。 """ import logging import time from pathlib import Path from typing import Any, Optional from app.plugins.run_log import get_run_log_service from app.utils.ocr_engine import get_ocr_engine, get_ocr_model_info logger = logging.getLogger(__name__) def parse_result(result) -> list[str]: """从 PP-OCRv6 predict 返回中抽取文本行。""" if not result: return [] lines: list[str] = [] for page in result: if not isinstance(page, dict): continue # 仅保留非空字符串 lines.extend(t for t in (page.get("rec_texts") or []) if t) return lines def _normalize_poly(poly) -> list | None: """把 PP-OCRv6 各种检测框格式统一为 4 点 polygon [[x,y], ...]。 rec_boxes 给的是 [xmin,ymin,xmax,ymax],polys/dt_polys 已经是 [[x,y], ...]。 """ if poly is None: return None arr = poly.tolist() if hasattr(poly, "tolist") else poly if not isinstance(arr, (list, tuple)) or len(arr) == 0: return None # 扁平 [xmin,ymin,xmax,ymax] 转四点 if len(arr) == 4 and all(isinstance(v, (int, float)) for v in arr): x1, y1, x2, y2 = arr return [[x1, y1], [x2, y1], [x2, y2], [x1, y2]] return [list(p) for p in arr] def extract_boxes(result) -> list[dict]: """抽取每行文本的检测框坐标与置信度,供前端叠加显示。""" boxes: list[dict] = [] if not result: return boxes for page in result: if not isinstance(page, dict): continue texts = page.get("rec_texts") or [] scores = page.get("rec_scores") or [] # rec_boxes 是矩形,其他字段已是 polygon polys = page.get("rec_polys") or page.get("dt_polys") or [] rects = page.get("rec_boxes") or [] for i, text in enumerate(texts): if not text: continue # 优先用 polygon;没有则退回到 rec_boxes 矩形 poly = _normalize_poly(polys[i] if i < len(polys) else None) if poly is None and i < len(rects): poly = _normalize_poly(rects[i]) boxes.append({ "text": str(text), "score": float(scores[i]) if i < len(scores) else None, "poly": poly, }) return boxes def extract_text(image_path: str, run_id: Optional[str] = None) -> dict: """从本地图片文件识别文字,可选地写 run 事件。 返回字段: success / text / lines / boxes / raw_result / model_info / timings / error """ run_service = get_run_log_service() if run_id else None def emit(stage: str, message: str, *, level: str = "info", detail: str | None = None) -> None: if run_service: run_service.add_event(run_id, stage, message, level=level, detail=detail) path = Path(image_path) if not path.exists(): msg = f"图片文件不存在: {image_path}" emit("ocr", msg, level="error") return _fail(msg, run_service) start = time.perf_counter() try: emit("model_load", "OCR 引擎已加载") engine = get_ocr_engine() emit("recognize", "开始识别") result = engine.predict(str(path)) lines = parse_result(result) boxes = extract_boxes(result) timings = {"total_ms": round((time.perf_counter() - start) * 1000, 2)} text = "\n".join(lines) emit("recognize", f"识别完成: {len(lines)} 行") return { "success": True, "text": text, "lines": lines, "boxes": boxes, "raw_result": result, "model_info": get_ocr_model_info(), "timings": timings, "error": None, } except Exception as e: logger.exception("OCR 识别失败: %s", image_path) emit("ocr", f"OCR 识别失败: {e}", level="error") return _fail(str(e), run_service) def _fail(error: str, run_service) -> dict: """构造统一的失败返回;为减少模板,只暴露必要字段。""" return { "success": False, "text": "", "lines": [], "boxes": [], "raw_result": None, "model_info": get_ocr_model_info() if run_service else None, "timings": None, "error": error, } def get_model_info_with_providers() -> dict: """聚合返回模型信息与 provider 可用性,前端一次拿全。""" from app.utils.ocr_engine import probe_ocr_providers return { "model_info": get_ocr_model_info(), "providers": {k: v.value for k, v in probe_ocr_providers().items()}, } # ---- 向后兼容层(被 content 调用)---- class OCRService: """薄壳:保持 `get_ocr_service().extract_text_from_file(path)` 旧接口可用。 内部直接调用模块级 extract_text,不再额外缓存引擎实例。 """ def extract_text_from_file(self, image_path: str) -> dict: return extract_text(image_path) def extract_text_from_base64(self, image_base64: str, suffix: str = ".jpg") -> dict: import base64 import tempfile try: content = base64.b64decode(image_base64) except Exception as e: return {"success": False, "text": "", "lines": [], "error": f"Base64 解码失败: {e}"} if suffix and not suffix.startswith("."): suffix = f".{suffix}" with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: tmp.write(content) tmp_path = tmp.name try: return self.extract_text_from_file(tmp_path) finally: Path(tmp_path).unlink(missing_ok=True) _ocr_service: OCRService | None = None def get_ocr_service() -> OCRService: """返回 OCRService 单例,供 content 等历史调用方使用。""" global _ocr_service if _ocr_service is None: _ocr_service = OCRService() return _ocr_service