guoju0808/ry-ocr/app/core/ocr_engine.py

"""PaddleOCR 引擎封装（单例 + 超时）

适配 paddleocr 3.0+ 的 predict() 接口。
降级到 2.x 时也能跑（参数兼容）。
"""
from __future__ import annotations

import threading
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutTimeout
from pathlib import Path
from typing import List, Optional

from loguru import logger

from app.config import settings
from app.models import OCRLine

_lock = threading.Lock()
_engine = None
_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="ocr")


# ---------- 引擎构造 ----------

def _build_engine():
    try:
        from paddleocr import PaddleOCR
        # paddleocr >= 3.0: device="cpu"/"gpu", 没有 use_gpu / show_log
        # 模型选择：
        #   v5_server - 精度高但 CPU 慢 (5+ 秒/页)
        #   v5_mobile - CPU 友好 (1~2 秒/页), 精度略低, 推荐生产环境
        mobile = settings.ocr_lang.endswith("mobile") or settings.ocr_engine == "mobile"
        if settings.use_gpu or settings.ocr_engine == "server":
            # GPU 或用户显式指定 server 模型
            det_name = "PP-OCRv5_server_det"
            rec_name = "PP-OCRv5_server_rec"
        else:
            # 默认 mobile (CPU 友好)
            det_name = "PP-OCRv5_mobile_det"
            rec_name = "PP-OCRv5_mobile_rec"

        return PaddleOCR(
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
            lang=settings.ocr_lang,
            device="gpu" if settings.use_gpu else "cpu",
            text_detection_model_name=det_name,
            text_recognition_model_name=rec_name,
        )
    except TypeError:
        # paddleocr 2.x 老参数（向后兼容）
        from paddleocr import PaddleOCR
        return PaddleOCR(
            use_angle_cls=True,
            lang=settings.ocr_lang,
            use_gpu=settings.use_gpu,
            show_log=False,
        )


def get_engine():
    global _engine
    if _engine is None:
        with _lock:
            if _engine is None:
                logger.info("正在初始化 PaddleOCR (lang={}, gpu={})...", settings.ocr_lang, settings.use_gpu)
                _engine = _build_engine()
                logger.info("PaddleOCR 初始化完成")
    return _engine


def _normalize(result) -> List[OCRLine]:
    lines: List[OCRLine] = []
    if not result:
        return lines

    if isinstance(result, list) and result and isinstance(result[0], dict):
        for page in result:
            texts = page.get("rec_texts") or []
            scores = page.get("rec_scores") or []
            polys = page.get("rec_polys") or []
            for i, txt in enumerate(texts):
                conf = float(scores[i]) if i < len(scores) else 0.0
                box = polys[i].tolist() if i < len(polys) and hasattr(polys[i], "tolist") else []
                lines.append(OCRLine(text=str(txt).strip(), confidence=conf, box=box))
        return lines

    try:
        for page in result:
            for det in page:
                box = det[0]
                txt, conf = det[1]
                lines.append(OCRLine(text=str(txt).strip(), confidence=float(conf), box=box))
    except Exception:
        logger.exception("无法解析 PaddleOCR 输出: {}", result)
    return lines


class OCRTimeout(Exception):
    """OCR 识别超时"""
    pass


def recognize(image_path: str | Path, timeout_s: Optional[float] = None) -> List[OCRLine]:
    """对单张图片做 OCR，带超时控制

    timeout_s 默认为 settings.ocr_page_timeout_s（默认 5 秒）。
    超时立即抛 OCRTimeout，不等模型跑完。
    """
    timeout = timeout_s or settings.ocr_page_timeout_s
    engine = get_engine()

    future = _executor.submit(engine.predict, str(image_path))
    try:
        raw = future.result(timeout=timeout)
    except FutTimeout:
        future.cancel()
        logger.error("OCR 超时 ({:.1f}s): {}", timeout, image_path)
        raise OCRTimeout(f"OCR 识别超时 ({timeout}秒): {image_path}")
    except Exception as e:
        logger.exception("OCR 执行异常: {}", e)
        raise

    return _normalize(raw)


def warmup():
    """预热：构造引擎 + 真跑一次空白图预测（避免首次请求超时）

    跳过预测失败（比如没装 paddle）但不报错。
    """
    try:
        engine = get_engine()
    except Exception as e:
        logger.warning("OCR 引擎初始化失败: {}", e)
        return

    try:
        import numpy as np
        # 1x1 灰度图，足以触发模型完整链路
        dummy = np.zeros((64, 64, 3), dtype=np.uint8)
        engine.predict(dummy)
        logger.info("OCR 预热完成（首次预测已跑）")
    except Exception as e:
        logger.warning("OCR 预测预热失败（不影响主流程）: {}", e)