"""Device validation and lazy PP-OCRv6 pipeline creation.""" from __future__ import annotations import logging import os import platform import time from dataclasses import dataclass from typing import Any from .result_adapter import DEFAULT_MODEL_SIZE, OCR_VERSION, model_names_for_size @dataclass class RuntimeConfig: device: str threads: int | None = None device_id: int = 0 lang: str = "ch" use_doc_orientation_classify: bool = True use_doc_unwarping: bool = False use_textline_orientation: bool = True text_recognition_batch_size: int = 6 model_size: str = DEFAULT_MODEL_SIZE @property def text_detection_model_name(self) -> str: return model_names_for_size(self.model_size)[0] @property def text_recognition_model_name(self) -> str: return model_names_for_size(self.model_size)[1] class PipelineProvider: def __init__(self, config: RuntimeConfig, logger: logging.Logger): self.config = config self.logger = logger self._pipeline: Any | None = None self._paddle: Any | None = None self._device_name: str | None = None self.import_seconds = 0.0 self.setup_seconds = 0.0 self.model_init_seconds = 0.0 @property def resolved_device(self) -> str: return "cpu" if self.config.device == "cpu" else f"gpu:{self.config.device_id}" def prepare(self) -> None: if self._paddle is not None: return started = time.perf_counter() try: import paddle except ImportError as exc: package = "paddlepaddle" if self.config.device == "cpu" else "paddlepaddle-gpu" raise RuntimeError(f"当前子项目未安装 {package}") from exc self.import_seconds = time.perf_counter() - started self._paddle = paddle setup_started = time.perf_counter() if self.config.device == "cpu": from paddle import core total_cores = os.cpu_count() or 4 threads = self.config.threads or max(1, total_cores - 2) if threads < 1: raise ValueError("CPU 线程数必须大于等于 1") self.config.threads = threads core.set_num_threads(threads) paddle.set_device("cpu") self._device_name = platform.processor() or "CPU" self.logger.info("CPU_CONFIGURED threads=%d total_cores=%d", threads, total_cores) else: if not paddle.is_compiled_with_cuda(): raise RuntimeError("当前 PaddlePaddle 不是 CUDA 构建；不会回退到 CPU") device_count = paddle.device.cuda.device_count() if device_count < 1: raise RuntimeError("未检测到 NVIDIA CUDA GPU；不会回退到 CPU") if self.config.device_id < 0 or self.config.device_id >= device_count: raise RuntimeError(f"GPU {self.config.device_id} 不存在，当前检测到 {device_count} 个设备") paddle.set_device(self.resolved_device) paddle.device.cuda.synchronize(self.config.device_id) try: self._device_name = paddle.device.cuda.get_device_name(self.config.device_id) except Exception: self._device_name = "unknown" self.logger.info("GPU_CONFIGURED device=%s device_name=%s", self.resolved_device, self._device_name) self.setup_seconds = time.perf_counter() - setup_started self.logger.info( "RUNTIME_PREPARED device=%s paddle_version=%s import_seconds=%.3f setup_seconds=%.3f", self.resolved_device, paddle.__version__, self.import_seconds, self.setup_seconds, ) def get(self): self.prepare() if self._pipeline is None: self.logger.info( "MODEL_INITIALIZATION_STARTED ocr_version=%s model_size=%s detection_model=%s recognition_model=%s device=%s", OCR_VERSION, self.config.model_size, self.config.text_detection_model_name, self.config.text_recognition_model_name, self.resolved_device, ) started = time.perf_counter() from paddleocr import PaddleOCR self._pipeline = PaddleOCR( text_detection_model_name=self.config.text_detection_model_name, text_recognition_model_name=self.config.text_recognition_model_name, device=self.resolved_device, use_doc_orientation_classify=self.config.use_doc_orientation_classify, use_doc_unwarping=self.config.use_doc_unwarping, use_textline_orientation=self.config.use_textline_orientation, text_recognition_batch_size=self.config.text_recognition_batch_size, ) self.synchronize() self.model_init_seconds = time.perf_counter() - started self.logger.info("MODEL_INITIALIZED seconds=%.3f device=%s", self.model_init_seconds, self.resolved_device) return self._pipeline def synchronize(self) -> None: if self.config.device == "gpu" and self._paddle is not None: self._paddle.device.cuda.synchronize(self.config.device_id) def gpu_memory(self) -> dict[str, float | None]: stats = {"allocated_mb": None, "reserved_mb": None, "max_allocated_mb": None, "max_reserved_mb": None} if self.config.device != "gpu" or self._paddle is None: return stats for key, name in { "allocated_mb": "memory_allocated", "reserved_mb": "memory_reserved", "max_allocated_mb": "max_memory_allocated", "max_reserved_mb": "max_memory_reserved", }.items(): function = getattr(self._paddle.device.cuda, name, None) if function is not None: try: stats[key] = round(float(function(self.config.device_id)) / (1024**2), 2) except Exception: pass return stats def model_config(self) -> dict[str, Any]: return { "ocr_version": OCR_VERSION, "model_size": self.config.model_size, "language": self.config.lang, "detection_model": self.config.text_detection_model_name, "recognition_model": self.config.text_recognition_model_name, "use_doc_orientation_classify": self.config.use_doc_orientation_classify, "use_doc_unwarping": self.config.use_doc_unwarping, "use_textline_orientation": self.config.use_textline_orientation, "text_recognition_batch_size": self.config.text_recognition_batch_size, } def metadata(self) -> dict[str, Any]: return { "device": self.resolved_device, "device_name": self._device_name, "cpu_threads": self.config.threads if self.config.device == "cpu" else None, "python_version": platform.python_version(), "platform": platform.platform(), "paddle_version": self._paddle.__version__ if self._paddle is not None else None, "runtime_import_seconds": round(self.import_seconds, 3), "runtime_setup_seconds": round(self.setup_seconds, 3), "model_init_seconds": round(self.model_init_seconds, 3), "model_config": self.model_config(), }