""" OCR 引擎模块 - 封装 PaddleOCR 单例 支持 PP-OCRv6 medium det/rec 成对配置,从 HuggingFace 加载权重。 """ import logging import os from typing import Optional, Dict, Any from pathlib import Path from app.plugins.model_probe import probe_provider, ProviderAvailability logger = logging.getLogger(__name__) # 全局 OCR 引擎实例 _ocr_engine: Optional[object] = None _ocr_model_info: Optional[Dict[str, Any]] = None # PP-OCRv6 模型配置 PP_OCRV6_MODELS = { "det": { "model_name": "PP-OCRv6_medium_det", "hf_model_id": "PaddlePaddle/PP-OCRv6_medium_det_safetensors", "provider": "ppocrv6", }, "rec": { "model_name": "PP-OCRv6_medium_rec", "hf_model_id": "PaddlePaddle/PP-OCRv6_medium_rec_safetensors", "provider": "ppocrv6_rec", }, } def get_ocr_model_config() -> Dict[str, Any]: """获取 OCR 模型配置 从环境变量或默认配置读取 det/rec 模型 ID。 Returns: Dict 包含 det 和 rec 的模型配置 """ det_model = os.getenv("OCR_DET_MODEL", PP_OCRV6_MODELS["det"]["model_name"]) rec_model = os.getenv("OCR_REC_MODEL", PP_OCRV6_MODELS["rec"]["model_name"]) return { "det": { "model_name": det_model, "hf_model_id": PP_OCRV6_MODELS["det"]["hf_model_id"], "provider": PP_OCRV6_MODELS["det"]["provider"], }, "rec": { "model_name": rec_model, "hf_model_id": PP_OCRV6_MODELS["rec"]["hf_model_id"], "provider": PP_OCRV6_MODELS["rec"]["provider"], }, } def probe_ocr_providers() -> Dict[str, ProviderAvailability]: """探测 OCR provider 的可用性 检查 paddleocr 和 paddle 依赖是否可用。 Returns: Dict 包含 det 和 rec provider 的可用性状态 """ config = get_ocr_model_config() results = {} for key, cfg in config.items(): status = probe_provider( provider=cfg["provider"], model_id=cfg["hf_model_id"], check_dependencies=["paddleocr", "paddle"], ) results[key] = status.status return results def get_ocr_engine(): """ 获取全局 OCR 引擎实例(单例模式) 使用 PP-OCRv6 medium det/rec 成对配置,从 HuggingFace 加载权重。 Returns: PaddleOCR 实例 """ global _ocr_engine, _ocr_model_info if _ocr_engine is None: config = get_ocr_model_config() det_cfg = config["det"] rec_cfg = config["rec"] logger.info(f"正在初始化 OCR 引擎: det={det_cfg['model_name']}, rec={rec_cfg['model_name']}") try: from paddleocr import PaddleOCR # 使用 PP-OCRv6 medium det/rec 成对配置 # 使用 transformers engine 避免 paddle inference 的兼容性问题 _ocr_engine = PaddleOCR( text_detection_model_name=det_cfg["model_name"], text_recognition_model_name=rec_cfg["model_name"], use_doc_orientation_classify=False, # 禁用文档方向分类 use_doc_unwarping=False, # 禁用文档矫正 use_textline_orientation=False, # 禁用文本行方向分类 lang='ch', # 中文模型 engine="transformers", # 使用 transformers engine ) # 记录模型信息 _ocr_model_info = { "det_model": det_cfg["model_name"], "rec_model": rec_cfg["model_name"], "det_hf_id": det_cfg["hf_model_id"], "rec_hf_id": rec_cfg["hf_model_id"], } logger.info("OCR 引擎初始化完成") except Exception as e: logger.error(f"OCR 引擎初始化失败: {e}") raise return _ocr_engine def get_ocr_model_info() -> Dict[str, Any]: """获取当前 OCR 模型信息 Returns: Dict 包含 det_model、rec_model 等信息 """ global _ocr_model_info if _ocr_model_info is None: config = get_ocr_model_config() _ocr_model_info = { "det_model": config["det"]["model_name"], "rec_model": config["rec"]["model_name"], "det_hf_id": config["det"]["hf_model_id"], "rec_hf_id": config["rec"]["hf_model_id"], } return _ocr_model_info def preload_ocr_engine(): """ 预加载 OCR 引擎 在应用启动时调用,提前加载模型以加快后续响应速度 """ try: engine = get_ocr_engine() logger.info("OCR 模型预加载成功") return engine except Exception as e: logger.error(f"OCR 模型预加载失败: {e}") raise