From 0fa7a9664c61de5a07d4fcb1a81e296f6764a691 Mon Sep 17 00:00:00 2001 From: SWHL Date: Wed, 8 Jul 2026 19:23:17 +0800 Subject: [PATCH] fix: fixed issue #695 --- python/rapidocr/default_models.yaml | 4 +- .../inference_engine/tensorrt/main.py | 96 +------------------ 2 files changed, 7 insertions(+), 93 deletions(-) diff --git a/python/rapidocr/default_models.yaml b/python/rapidocr/default_models.yaml index a5474dc..067e6b3 100644 --- a/python/rapidocr/default_models.yaml +++ b/python/rapidocr/default_models.yaml @@ -11,8 +11,8 @@ onnxruntime: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.1/onnx/PP-OCRv4/det/en_PP-OCRv3_det_mobile.onnx SHA256: ea07c15d38ac40cd69da3c493444ec75b44ff23840553ff8ba102c1219ed39c2 multi_PP-OCRv3_det_mobile: - model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.1/onnx/PP-OCRv4/det/multi_PP-OCRv3_det_mobile.onnx - SHA256: 5475c6c7f4d84a6c4f32241b487435d59f126a40c023387af99732258844cdc3 + model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/master/onnx/PP-OCRv4/det/multi_PP-OCRv3_det_mobile.onnx + SHA256: 5582933662898cd2b4c2eae8a7e5e4fbdff84773711f29db00a1cbd185059001 rec: arabic_PP-OCRv4_rec_mobile: model_dir: https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.1/onnx/PP-OCRv4/rec/arabic_PP-OCRv4_rec_mobile.onnx diff --git a/python/rapidocr/inference_engine/tensorrt/main.py b/python/rapidocr/inference_engine/tensorrt/main.py index b0d1392..5968cdf 100644 --- a/python/rapidocr/inference_engine/tensorrt/main.py +++ b/python/rapidocr/inference_engine/tensorrt/main.py @@ -41,16 +41,7 @@ class TRTInferSession(InferSession): 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 + self._max_square_size = 2048 def __enter__(self) -> "TRTInferSession": return self @@ -114,27 +105,18 @@ class TRTInferSession(InferSession): >>> 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 + # Step 1: Set input shape for dynamic dimensions output_shape = self._set_input_shape(input_content) - # Step 3: Copy input data to GPU + # Step 2: Copy input data to GPU self._copy_input_to_device(input_content) - # Step 4: Execute inference + # Step 3: Execute inference self._execute_inference() - # Step 5: Copy output back to CPU + # Step 4: 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: @@ -183,16 +165,6 @@ class TRTInferSession(InferSession): 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 @@ -213,64 +185,6 @@ class TRTInferSession(InferSession): # 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)