# -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com 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.engine.create_execution_context() # 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}") # Detect MULTI model for square input requirement self._requires_square_input = self._check_multi_model(cfg) if self._requires_square_input: self._max_square_size = self._get_max_profile_size() logger.debug( f"MULTI det model: requires square input, max_size={self._max_square_size}" ) else: self._requires_square_input = False 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: Handle square input for MULTI model original_hw = None if self._requires_square_input: input_content, original_hw = self._pad_to_square(input_content) # Step 2: Set input shape for dynamic dimensions output_shape = self._set_input_shape(input_content) # Step 3: Copy input data to GPU self._copy_input_to_device(input_content) # Step 4: Execute inference self._execute_inference() # Step 5: Copy output back to CPU output = self._copy_output_to_host(output_shape) # Step 6: Crop output back to original shape if original_hw is not None: output = self._crop_output(output, original_hw) 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 _check_multi_model(self, cfg: Dict[str, Any]) -> bool: try: from ...utils.typings import LangDet, TaskType is_det = getattr(cfg, "task_type", None) == TaskType.DET is_multi = getattr(cfg, "lang_type", None) == LangDet.MULTI return is_det and is_multi except Exception: return False 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 _pad_to_square(self, input_content: np.ndarray) -> tuple: N, C, H, W = input_content.shape if H == W: return input_content, None # Already square # Calculate square size square_size = max(H, W) # Limit to max profile size if needed if square_size > self._max_square_size: logger.warning( f"Square size {square_size} exceeds max profile size {self._max_square_size}. " f"Limiting to {self._max_square_size}. This may cause accuracy loss." ) square_size = self._max_square_size # Ensure divisible by 32 (TensorRT requirement) square_size = int(round(square_size / 32) * 32) # Create padded array (zero padding) padded = np.zeros((N, C, square_size, square_size), dtype=input_content.dtype) # Copy original content to top-left copy_h = min(H, square_size) copy_w = min(W, square_size) padded[:, :, :copy_h, :copy_w] = input_content[:, :, :copy_h, :copy_w] return padded, (H, W) def _crop_output(self, output: np.ndarray, original_hw: tuple) -> np.ndarray: """Crop output back to original shape after square inference. Args: output: Output array with shape (N, C, S, S) from square inference. original_hw: Original (H, W) before padding. Returns: Cropped output array with shape (N, C, H, W). """ if original_hw is None: return output orig_h, orig_w = original_hw out_h, out_w = output.shape[2:4] if out_h >= orig_h and out_w >= orig_w: # Normal case: crop to original size return output[:, :, :orig_h, :orig_w] else: # Output is smaller - need to scale # This happens when square_size was limited by max_profile_size scale_h = out_h / max(orig_h, orig_w) scale_w = out_w / max(orig_h, orig_w) crop_h = int(orig_h * scale_h) crop_w = int(orig_w * scale_w) return output[:, :, :crop_h, :crop_w] 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" return cache_dir / f"{model_name}_{gpu_arch}_{precision}.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() return runtime.deserialize_cuda_engine(engine_data) 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