mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 21:10:55 +08:00
386 lines
14 KiB
Python
386 lines
14 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
import os
|
|
import traceback
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import numpy as np
|
|
import tensorrt as trt
|
|
from cuda.bindings import runtime as cudart
|
|
|
|
from ...utils.download_file import DownloadFile, DownloadFileInput
|
|
from ...utils.log import logger
|
|
from ...utils.model_resolver import normalize_lang, resolve_model_key
|
|
from ...utils.typings import EngineType
|
|
from ...utils.utils import mkdir
|
|
from ..base import FileInfo, InferSession
|
|
from .engine_builder import TRTEngineBuilder
|
|
from .memory_utils import allocate_buffers, free_buffers
|
|
|
|
|
|
class TRTInferSession(InferSession):
|
|
def __init__(self, cfg: Dict[str, Any]):
|
|
self.cfg = cfg
|
|
self.engine_cfg = cfg.get("engine_cfg", {})
|
|
self.model_root_dir = None
|
|
self._closed = False
|
|
self.device_id = self._setup_cuda_device()
|
|
|
|
self.trt_logger = trt.Logger(trt.Logger.WARNING)
|
|
|
|
engine_path = self._get_engine_path(cfg)
|
|
self.engine = self._load_or_build_engine(cfg, engine_path)
|
|
|
|
self.context = self._create_execution_context(engine_path)
|
|
|
|
# Allocate memory buffers (pre-allocated with max shape)
|
|
self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(
|
|
self.engine, self.context
|
|
)
|
|
|
|
logger.info(f"TensorRT engine loaded: {engine_path}")
|
|
|
|
self._max_square_size = 2048
|
|
|
|
def __enter__(self) -> "TRTInferSession":
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
self.close()
|
|
|
|
def close(self) -> None:
|
|
if self._closed:
|
|
return
|
|
|
|
self._closed = True
|
|
|
|
try:
|
|
# Synchronize stream before cleanup
|
|
if hasattr(self, "stream") and self.stream is not None:
|
|
try:
|
|
cudart.cudaStreamSynchronize(self.stream)
|
|
except Exception as e:
|
|
logger.debug(f"Stream sync error during close: {e}")
|
|
|
|
# Free GPU memory buffers
|
|
if hasattr(self, "inputs") and hasattr(self, "outputs"):
|
|
try:
|
|
free_buffers(self.inputs, self.outputs, self.stream)
|
|
except Exception as e:
|
|
logger.debug(f"Buffer free error during close: {e}")
|
|
|
|
# Destroy CUDA stream
|
|
if hasattr(self, "stream") and self.stream is not None:
|
|
try:
|
|
cudart.cudaStreamDestroy(self.stream)
|
|
except Exception as e:
|
|
logger.debug(f"Stream destroy error during close: {e}")
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Error during session close: {e}")
|
|
|
|
def __del__(self):
|
|
self.close()
|
|
|
|
def __call__(self, input_content: np.ndarray) -> np.ndarray:
|
|
"""Run inference on input data.
|
|
|
|
This method executes the TensorRT engine on the provided input.
|
|
It uses pre-allocated buffers to avoid memory allocation overhead.
|
|
|
|
Args:
|
|
input_content: Input numpy array with shape matching the model.
|
|
For detection: (N, C, H, W)
|
|
For recognition: (N, C, H, W)
|
|
|
|
Returns:
|
|
Output numpy array from the model.
|
|
|
|
Raises:
|
|
TensorRTError: If inference fails.
|
|
|
|
Example:
|
|
>>> input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
|
|
>>> output = session(input_data)
|
|
"""
|
|
try:
|
|
# Step 1: Set input shape for dynamic dimensions
|
|
output_shape = self._set_input_shape(input_content)
|
|
|
|
# Step 2: Copy input data to GPU
|
|
self._copy_input_to_device(input_content)
|
|
|
|
# Step 3: Execute inference
|
|
self._execute_inference()
|
|
|
|
# Step 4: Copy output back to CPU
|
|
output = self._copy_output_to_host(output_shape)
|
|
|
|
return output
|
|
|
|
except Exception as e:
|
|
error_info = traceback.format_exc()
|
|
raise TensorRTError(
|
|
f"Inference failed for input shape {input_content.shape}:\n{error_info}"
|
|
) from e
|
|
|
|
def _set_input_shape(self, input_content: np.ndarray) -> tuple:
|
|
input_name = self.engine.get_tensor_name(0)
|
|
self.context.set_input_shape(input_name, input_content.shape)
|
|
|
|
output_name = self.engine.get_tensor_name(1)
|
|
return self.context.get_tensor_shape(output_name)
|
|
|
|
def _copy_input_to_device(self, input_content: np.ndarray) -> None:
|
|
input_flat = input_content.ravel()
|
|
self.inputs[0].host[: input_flat.size] = input_flat
|
|
|
|
cudart.cudaMemcpyAsync(
|
|
self.inputs[0].device,
|
|
self.inputs[0].host.ctypes.data,
|
|
input_flat.nbytes,
|
|
cudart.cudaMemcpyKind.cudaMemcpyHostToDevice,
|
|
self.stream,
|
|
)
|
|
|
|
def _execute_inference(self) -> None:
|
|
"""Execute TensorRT inference asynchronously."""
|
|
self.context.execute_async_v3(stream_handle=self.stream)
|
|
|
|
def _copy_output_to_host(self, output_shape: tuple) -> np.ndarray:
|
|
output_size = int(np.prod(output_shape))
|
|
output_nbytes = output_size * self.outputs[0].host.itemsize
|
|
|
|
cudart.cudaMemcpyAsync(
|
|
self.outputs[0].host.ctypes.data,
|
|
self.outputs[0].device,
|
|
output_nbytes,
|
|
cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost,
|
|
self.stream,
|
|
)
|
|
|
|
# Wait for all async operations to complete
|
|
cudart.cudaStreamSynchronize(self.stream)
|
|
|
|
return self.outputs[0].host[:output_size].reshape(output_shape)
|
|
|
|
def _get_max_profile_size(self) -> int:
|
|
try:
|
|
# Try to get from config first
|
|
profile_cfg = self.engine_cfg.get("det_profile", {})
|
|
max_shape = profile_cfg.get("max_shape")
|
|
if max_shape and len(max_shape) >= 4:
|
|
return max(max_shape[2], max_shape[3]) # max(H, W)
|
|
|
|
# Fallback: query from engine
|
|
input_name = self.engine.get_tensor_name(0)
|
|
profile_shape = self.engine.get_tensor_profile_shape(input_name, 0)
|
|
if profile_shape and len(profile_shape) >= 3:
|
|
max_shape = profile_shape[2] # Index 2 = max shape
|
|
return max(max_shape[2], max_shape[3])
|
|
except Exception as e:
|
|
logger.debug(f"Could not get max profile size: {e}")
|
|
|
|
# Default fallback
|
|
return 2048
|
|
|
|
def _setup_cuda_device(self) -> int:
|
|
device_id = self.engine_cfg.get("device_id", 0)
|
|
status_tuple = cudart.cudaSetDevice(device_id)
|
|
status = status_tuple[0]
|
|
assert status.value == 0, (
|
|
f"Failed to set CUDA device {device_id}: {status}. "
|
|
f"Ensure CUDA is properly installed and GPU is available."
|
|
)
|
|
return device_id
|
|
|
|
def _get_engine_path(self, cfg: Dict[str, Any]) -> Path:
|
|
cache_dir = self.engine_cfg.get("cache_dir")
|
|
if cache_dir is None:
|
|
# Check model_root_dir only if cache_dir is None
|
|
if self.model_root_dir is None:
|
|
model_root_dir = cfg.get("model_root_dir", None)
|
|
if model_root_dir is None:
|
|
raise ValueError(
|
|
"Either model_path or model_root_dir must be provided in the configuration."
|
|
)
|
|
model_root_dir = Path(model_root_dir)
|
|
mkdir(model_root_dir)
|
|
|
|
self.model_root_dir = model_root_dir
|
|
if not self.model_root_dir.exists():
|
|
raise FileNotFoundError(
|
|
f"model_root_dir {self.model_root_dir} does not exist"
|
|
)
|
|
|
|
cache_dir = self.model_root_dir / "models"
|
|
|
|
cache_dir = Path(cache_dir)
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
model_name = self._get_model_name(cfg)
|
|
gpu_arch = self._get_gpu_arch()
|
|
precision = "fp16" if self.engine_cfg.get("use_fp16", True) else "fp32"
|
|
tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset")
|
|
|
|
return cache_dir / (
|
|
f"{model_name}_{gpu_arch}_{precision}_tf32{tf32_override}.engine"
|
|
)
|
|
|
|
def _get_model_name(self, cfg: Dict[str, Any]) -> str:
|
|
"""Extract model name from config for engine filename."""
|
|
if cfg.get("model_path"):
|
|
return Path(cfg["model_path"]).stem
|
|
|
|
model_key = resolve_model_key(
|
|
cfg.task_type,
|
|
cfg.ocr_version,
|
|
cfg.lang_type,
|
|
cfg.model_type,
|
|
)
|
|
if model_key is not None:
|
|
return model_key
|
|
|
|
task_type = cfg.task_type.value
|
|
lang_type = normalize_lang(cfg.lang_type)
|
|
ocr_version = cfg.ocr_version.value
|
|
model_type = cfg.model_type.value
|
|
return f"{lang_type}_{ocr_version}_{task_type}_{model_type}"
|
|
|
|
def _get_gpu_arch(self) -> str:
|
|
"""Get GPU architecture string for cache key (e.g., 'sm87')."""
|
|
status, major = cudart.cudaDeviceGetAttribute(
|
|
cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, self.device_id
|
|
)
|
|
assert status.value == 0, f"Failed to get compute capability: {status}"
|
|
|
|
status, minor = cudart.cudaDeviceGetAttribute(
|
|
cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, self.device_id
|
|
)
|
|
assert status.value == 0, f"Failed to get compute capability: {status}"
|
|
|
|
return f"sm{major}{minor}"
|
|
|
|
def _load_or_build_engine(
|
|
self, cfg: Dict[str, Any], engine_path: Path
|
|
) -> trt.ICudaEngine:
|
|
force_rebuild = self.engine_cfg.get("force_rebuild", False)
|
|
|
|
# Try to load cached engine
|
|
if engine_path.exists() and not force_rebuild:
|
|
try:
|
|
return self._load_engine(engine_path)
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Failed to load cached engine {engine_path}: {e}. "
|
|
f"Will rebuild from ONNX."
|
|
)
|
|
|
|
# Build new engine from ONNX
|
|
onnx_path = self._get_onnx_path(cfg)
|
|
logger.info(f"Building TensorRT engine from {onnx_path}")
|
|
|
|
builder = TRTEngineBuilder(
|
|
onnx_path=onnx_path,
|
|
engine_path=engine_path,
|
|
cfg=self.engine_cfg,
|
|
task_type=cfg.task_type.value,
|
|
trt_logger=self.trt_logger,
|
|
ocr_version=cfg.ocr_version,
|
|
)
|
|
return builder.build()
|
|
|
|
def _get_onnx_path(self, cfg: Dict[str, Any]) -> Path:
|
|
model_path = cfg.get("model_path")
|
|
|
|
if model_path is None:
|
|
original_engine_type = cfg.engine_type
|
|
cfg.engine_type = EngineType.ONNXRUNTIME
|
|
|
|
model_info = self.get_model_url(
|
|
FileInfo(
|
|
engine_type=EngineType.ONNXRUNTIME,
|
|
ocr_version=cfg.ocr_version,
|
|
task_type=cfg.task_type,
|
|
lang_type=cfg.lang_type,
|
|
model_type=cfg.model_type,
|
|
)
|
|
)
|
|
|
|
cfg.engine_type = original_engine_type
|
|
|
|
# Check model_root_dir only if model_path is None
|
|
if self.model_root_dir is None:
|
|
self.model_root_dir = Path(cfg.get("model_root_dir"))
|
|
if not self.model_root_dir.exists():
|
|
raise FileNotFoundError(
|
|
f"model_root_dir {self.model_root_dir} does not exist"
|
|
)
|
|
|
|
model_path = self.model_root_dir / Path(model_info["model_dir"]).name
|
|
download_params = DownloadFileInput(
|
|
file_url=model_info["model_dir"],
|
|
sha256=model_info["SHA256"],
|
|
save_path=model_path,
|
|
logger=logger,
|
|
)
|
|
DownloadFile.run(download_params)
|
|
|
|
model_path = Path(model_path)
|
|
self._verify_model(model_path)
|
|
return model_path
|
|
|
|
def _load_engine(self, engine_path: Path) -> trt.ICudaEngine:
|
|
runtime = trt.Runtime(self.trt_logger)
|
|
with open(engine_path, "rb") as f:
|
|
engine_data = f.read()
|
|
|
|
engine = runtime.deserialize_cuda_engine(engine_data)
|
|
if engine is None:
|
|
raise TensorRTError(f"Failed to deserialize TensorRT engine: {engine_path}")
|
|
return engine
|
|
|
|
def _create_execution_context(self, engine_path: Path) -> trt.IExecutionContext:
|
|
context = self.engine.create_execution_context()
|
|
if context is None:
|
|
tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset")
|
|
raise TensorRTError(
|
|
"Failed to create TensorRT execution context for "
|
|
f"{engine_path}. NVIDIA_TF32_OVERRIDE={tf32_override}. "
|
|
"Remove the cached engine or rebuild it under the current "
|
|
"TensorRT/CUDA environment."
|
|
)
|
|
return context
|
|
|
|
def have_key(self, key: str = "character") -> bool:
|
|
return False
|
|
|
|
def get_character_list(self, key: str = "character") -> List[str]:
|
|
return []
|
|
|
|
@classmethod
|
|
def get_dict_key_url(cls, file_info: FileInfo) -> Optional[str]:
|
|
# Try Paddle first (usually has dict_url)
|
|
for engine_type in [EngineType.PADDLE, EngineType.ONNXRUNTIME]:
|
|
try:
|
|
fallback_info = FileInfo(
|
|
engine_type=engine_type,
|
|
ocr_version=file_info.ocr_version,
|
|
task_type=file_info.task_type,
|
|
lang_type=file_info.lang_type,
|
|
model_type=file_info.model_type,
|
|
)
|
|
model_dict = cls.get_model_url(fallback_info)
|
|
if model_dict and "dict_url" in model_dict:
|
|
return model_dict["dict_url"]
|
|
except Exception as e:
|
|
logger.debug(f"Failed to get dict URL from {engine_type.value}: {e}")
|
|
|
|
return None
|
|
|
|
|
|
class TensorRTError(Exception):
|
|
pass
|