176 lines
7.3 KiB
Python
176 lines
7.3 KiB
Python
"""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(),
|
||
}
|