fix: fixed issue #694

python/rapidocr/inference_engine/tensorrt/main.py

@@ -21,8 +21,8 @@ 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)

python/tests/test_engine_config.py

@@ -0,0 +1,206 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import sys
+from pathlib import Path
+
+import pytest
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+
+from rapidocr.inference_engine.onnxruntime import provider_config
+from rapidocr.inference_engine.onnxruntime.provider_config import EP, ProviderConfig
+from rapidocr.inference_engine.openvino.device_config import CPUConfig
+
+
+class AttrDict(dict):
+    def __getattr__(self, key):
+        try:
+            return self[key]
+        except KeyError as e:
+            raise AttributeError(key) from e
+
+    def __setattr__(self, key, value):
+        self[key] = value
+
+
+def make_provider_cfg(**overrides):
+    cfg = AttrDict(
+        use_cuda=False,
+        cuda_ep_cfg={"device_id": 0},
+        use_dml=False,
+        dml_ep_cfg=None,
+        use_cann=False,
+        cann_ep_cfg={"device_id": 0},
+        use_coreml=False,
+        coreml_ep_cfg={"ModelFormat": "MLProgram"},
+        cpu_ep_cfg={"arena_extend_strategy": "kSameAsRequested"},
+    )
+    cfg.update(overrides)
+    return cfg
+
+
+def make_provider_config(monkeypatch, cfg, providers, device="CPU"):
+    monkeypatch.setattr(provider_config, "get_available_providers", lambda: providers)
+    monkeypatch.setattr(provider_config, "get_device", lambda: device)
+    return ProviderConfig(cfg)
+
+
+def test_onnxruntime_provider_config_uses_cpu_by_default(monkeypatch):
+    cfg = make_provider_cfg()
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.CPU_EP.value],
+    )
+
+    assert provider_cfg.get_ep_list() == [(EP.CPU_EP.value, cfg.cpu_ep_cfg)]
+
+
+def test_onnxruntime_provider_config_prefers_cuda_when_available(monkeypatch):
+    cfg = make_provider_cfg(use_cuda=True)
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.CUDA_EP.value, EP.CPU_EP.value],
+        device="GPU",
+    )
+
+    assert provider_cfg.get_ep_list() == [
+        (EP.CUDA_EP.value, cfg.cuda_ep_cfg),
+        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
+    ]
+
+
+def test_onnxruntime_provider_config_falls_back_when_cuda_unavailable(monkeypatch):
+    cfg = make_provider_cfg(use_cuda=True)
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.CPU_EP.value],
+        device="CPU",
+    )
+
+    assert provider_cfg.get_ep_list() == [(EP.CPU_EP.value, cfg.cpu_ep_cfg)]
+
+
+def test_onnxruntime_provider_config_prefers_directml_on_windows(monkeypatch):
+    cfg = make_provider_cfg(use_dml=True, dml_ep_cfg={"device_id": 0})
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.DIRECTML_EP.value, EP.CPU_EP.value],
+    )
+    monkeypatch.setattr(provider_config.platform, "system", lambda: "Windows")
+    monkeypatch.setattr(provider_config.platform, "version", lambda: "10.0.19045")
+
+    assert provider_cfg.get_ep_list() == [
+        (EP.DIRECTML_EP.value, cfg.dml_ep_cfg),
+        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
+    ]
+
+
+def test_onnxruntime_provider_config_skips_directml_on_non_windows(monkeypatch):
+    cfg = make_provider_cfg(use_dml=True)
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.DIRECTML_EP.value, EP.CPU_EP.value],
+    )
+    monkeypatch.setattr(provider_config.platform, "system", lambda: "Linux")
+
+    assert provider_cfg.get_ep_list() == [(EP.CPU_EP.value, cfg.cpu_ep_cfg)]
+
+
+def test_onnxruntime_provider_config_adds_cann_when_available(monkeypatch):
+    cfg = make_provider_cfg(use_cann=True)
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.CANN_EP.value, EP.CPU_EP.value],
+    )
+
+    assert provider_cfg.get_ep_list() == [
+        (EP.CANN_EP.value, cfg.cann_ep_cfg),
+        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
+    ]
+
+
+def test_onnxruntime_provider_config_adds_coreml_on_darwin(monkeypatch):
+    cfg = make_provider_cfg(use_coreml=True)
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.COREML_EP.value, EP.CPU_EP.value],
+    )
+    monkeypatch.setattr(provider_config.platform, "system", lambda: "Darwin")
+
+    assert provider_cfg.get_ep_list() == [
+        (EP.COREML_EP.value, cfg.coreml_ep_cfg),
+        (EP.CPU_EP.value, cfg.cpu_ep_cfg),
+    ]
+
+
+def test_onnxruntime_provider_config_rejects_empty_session_providers(monkeypatch):
+    cfg = make_provider_cfg()
+    provider_cfg = make_provider_config(
+        monkeypatch,
+        cfg,
+        providers=[EP.CPU_EP.value],
+    )
+
+    with pytest.raises(ValueError, match="Session Providers is empty"):
+        provider_cfg.verify_providers([])
+
+
+def test_openvino_cpu_config_omits_defaults():
+    cfg = AttrDict(
+        inference_num_threads=-1,
+        performance_hint=None,
+        performance_num_requests=-1,
+        enable_cpu_pinning=None,
+        num_streams=-1,
+        enable_hyper_threading=None,
+        scheduling_core_type=None,
+    )
+
+    assert CPUConfig(cfg).get_config() == {}
+
+
+def test_openvino_cpu_config_maps_non_default_values(monkeypatch):
+    cfg = AttrDict(
+        inference_num_threads=4,
+        performance_hint="LATENCY",
+        performance_num_requests=2,
+        enable_cpu_pinning=True,
+        num_streams=1,
+        enable_hyper_threading=False,
+        scheduling_core_type="PCORE_ONLY",
+    )
+    monkeypatch.setattr("rapidocr.inference_engine.openvino.device_config.os.cpu_count", lambda: 8)
+
+    assert CPUConfig(cfg).get_config() == {
+        "INFERENCE_NUM_THREADS": "4",
+        "PERFORMANCE_HINT": "LATENCY",
+        "PERFORMANCE_HINT_NUM_REQUESTS": "2",
+        "ENABLE_CPU_PINNING": "True",
+        "NUM_STREAMS": "1",
+        "ENABLE_HYPER_THREADING": "False",
+        "SCHEDULING_CORE_TYPE": "PCORE_ONLY",
+    }
+
+
+def test_openvino_cpu_config_ignores_invalid_thread_count(monkeypatch):
+    cfg = AttrDict(
+        inference_num_threads=16,
+        performance_hint=None,
+        performance_num_requests=-1,
+        enable_cpu_pinning=None,
+        num_streams=-1,
+        enable_hyper_threading=None,
+        scheduling_core_type=None,
+    )
+    monkeypatch.setattr("rapidocr.inference_engine.openvino.device_config.os.cpu_count", lambda: 8)
+
+    assert "INFERENCE_NUM_THREADS" not in CPUConfig(cfg).get_config()

python/tests/test_tensorrt.py

@@ -0,0 +1,345 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import sys
+import types
+from pathlib import Path
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+
+from rapidocr.utils.typings import EngineType, LangDet, ModelType, OCRVersion, TaskType
+
+
+class AttrDict(dict):
+    def __getattr__(self, key):
+        try:
+            return self[key]
+        except KeyError as e:
+            raise AttributeError(key) from e
+
+    def __setattr__(self, key, value):
+        self[key] = value
+
+
+def install_fake_tensorrt_modules(monkeypatch):
+    trt_module = types.ModuleType("tensorrt")
+
+    class Logger:
+        WARNING = 2
+
+        def __init__(self, severity=None):
+            self.severity = severity
+
+    class TensorIOMode:
+        INPUT = "input"
+        OUTPUT = "output"
+
+    trt_module.Logger = Logger
+    trt_module.ICudaEngine = object
+    trt_module.IExecutionContext = object
+    trt_module.INetworkDefinition = object
+    trt_module.IOptimizationProfile = object
+    trt_module.TensorIOMode = TensorIOMode
+    trt_module.nptype = lambda dtype: dtype
+    trt_module.volume = lambda shape: 1
+
+    cuda_module = types.ModuleType("cuda")
+    bindings_module = types.ModuleType("cuda.bindings")
+    runtime_module = types.ModuleType("cuda.bindings.runtime")
+
+    class CudaStatus:
+        value = 0
+
+    class CudaDeviceAttr:
+        cudaDevAttrComputeCapabilityMajor = "major"
+        cudaDevAttrComputeCapabilityMinor = "minor"
+
+    class CudaMemcpyKind:
+        cudaMemcpyHostToDevice = "host_to_device"
+        cudaMemcpyDeviceToHost = "device_to_host"
+
+    runtime_module.cudaDeviceAttr = CudaDeviceAttr
+    runtime_module.cudaMemcpyKind = CudaMemcpyKind
+    runtime_module.cudaSetDevice = lambda device_id: (CudaStatus(),)
+    runtime_module.cudaDeviceGetAttribute = lambda attr, device_id: (
+        CudaStatus(),
+        8 if attr == CudaDeviceAttr.cudaDevAttrComputeCapabilityMajor else 7,
+    )
+    runtime_module.cudaStreamSynchronize = lambda stream: (CudaStatus(),)
+    runtime_module.cudaStreamDestroy = lambda stream: (CudaStatus(),)
+
+    monkeypatch.setitem(sys.modules, "tensorrt", trt_module)
+    monkeypatch.setitem(sys.modules, "cuda", cuda_module)
+    monkeypatch.setitem(sys.modules, "cuda.bindings", bindings_module)
+    monkeypatch.setitem(sys.modules, "cuda.bindings.runtime", runtime_module)
+
+
+class FakeEngine:
+    def create_execution_context(self):
+        return object()
+
+
+def make_tensorrt_cfg(tmp_path, **overrides):
+    model_root_dir = tmp_path / "models"
+    model_root_dir.mkdir(exist_ok=True)
+
+    cfg = AttrDict(
+        engine_type=EngineType.TENSORRT,
+        engine_cfg={},
+        model_root_dir=model_root_dir,
+        model_path=None,
+        ocr_version=OCRVersion.PPOCRV4,
+        task_type=TaskType.DET,
+        lang_type=LangDet.EN,
+        model_type=ModelType.MOBILE,
+    )
+    cfg.update(overrides)
+    return cfg
+
+
+def import_tensorrt_main(monkeypatch):
+    install_fake_tensorrt_modules(monkeypatch)
+
+    from rapidocr.inference_engine.tensorrt import main as tensorrt_main
+
+    monkeypatch.setattr(
+        tensorrt_main,
+        "allocate_buffers",
+        lambda engine, context: ([], [], [], None),
+    )
+    return tensorrt_main
+
+
+def test_tensorrt_cache_dir_initializes_model_root_dir_for_build(
+    monkeypatch, tmp_path
+):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    downloaded_paths = []
+
+    def fake_download(input_params):
+        save_path = Path(input_params.save_path)
+        save_path.write_bytes(b"fake onnx model")
+        downloaded_paths.append(save_path)
+
+    class FakeEngineBuilder:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+        def build(self):
+            return FakeEngine()
+
+    monkeypatch.setattr(tensorrt_main.DownloadFile, "run", fake_download)
+    monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", FakeEngineBuilder)
+
+    cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"})
+
+    session = tensorrt_main.TRTInferSession(cfg)
+
+    assert session.model_root_dir == cfg.model_root_dir
+    assert downloaded_paths == [cfg.model_root_dir / "en_PP-OCRv3_det_mobile.onnx"]
+
+
+def test_tensorrt_default_cache_dir_uses_model_root_dir(monkeypatch, tmp_path):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    engine_paths = []
+
+    class FakeEngineBuilder:
+        def __init__(self, **kwargs):
+            engine_paths.append(kwargs["engine_path"])
+
+        def build(self):
+            return FakeEngine()
+
+    monkeypatch.setattr(
+        tensorrt_main.DownloadFile,
+        "run",
+        lambda input_params: Path(input_params.save_path).write_bytes(b"fake onnx"),
+    )
+    monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", FakeEngineBuilder)
+
+    cfg = make_tensorrt_cfg(tmp_path)
+
+    session = tensorrt_main.TRTInferSession(cfg)
+
+    assert session.model_root_dir == cfg.model_root_dir
+    assert engine_paths == [
+        cfg.model_root_dir / "models" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine"
+    ]
+
+
+def test_tensorrt_explicit_model_path_skips_download(monkeypatch, tmp_path):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    model_path = tmp_path / "custom.onnx"
+    model_path.write_bytes(b"fake onnx")
+    builder_kwargs = []
+
+    class FakeEngineBuilder:
+        def __init__(self, **kwargs):
+            builder_kwargs.append(kwargs)
+
+        def build(self):
+            return FakeEngine()
+
+    def fail_download(input_params):
+        raise AssertionError("explicit model_path should not download a default model")
+
+    monkeypatch.setattr(tensorrt_main.DownloadFile, "run", fail_download)
+    monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", FakeEngineBuilder)
+
+    cfg = make_tensorrt_cfg(
+        tmp_path,
+        engine_cfg={"cache_dir": tmp_path / "trt_cache", "use_fp16": False},
+        model_path=model_path,
+    )
+
+    tensorrt_main.TRTInferSession(cfg)
+
+    assert builder_kwargs[0]["onnx_path"] == model_path
+    assert builder_kwargs[0]["engine_path"] == (
+        tmp_path / "trt_cache" / "custom_sm87_fp32.engine"
+    )
+
+
+def test_tensorrt_loads_cached_engine_without_rebuild(monkeypatch, tmp_path):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"})
+    engine_path = tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine"
+    engine_path.parent.mkdir()
+    engine_path.write_bytes(b"cached engine")
+
+    loaded_paths = []
+
+    def fake_load_engine(self, path):
+        loaded_paths.append(path)
+        return FakeEngine()
+
+    def fail_build(*args, **kwargs):
+        raise AssertionError("cached engine should be loaded without rebuilding")
+
+    monkeypatch.setattr(tensorrt_main.TRTInferSession, "_load_engine", fake_load_engine)
+    monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", fail_build)
+
+    tensorrt_main.TRTInferSession(cfg)
+
+    assert loaded_paths == [engine_path]
+
+
+def test_tensorrt_force_rebuild_ignores_cached_engine(monkeypatch, tmp_path):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    cfg = make_tensorrt_cfg(
+        tmp_path,
+        engine_cfg={"cache_dir": tmp_path / "trt_cache", "force_rebuild": True},
+    )
+    engine_path = tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine"
+    engine_path.parent.mkdir()
+    engine_path.write_bytes(b"cached engine")
+
+    build_count = 0
+
+    class FakeEngineBuilder:
+        def __init__(self, **kwargs):
+            pass
+
+        def build(self):
+            nonlocal build_count
+            build_count += 1
+            return FakeEngine()
+
+    def fail_load_engine(self, path):
+        raise AssertionError("force_rebuild should skip cached engine loading")
+
+    monkeypatch.setattr(
+        tensorrt_main.DownloadFile,
+        "run",
+        lambda input_params: Path(input_params.save_path).write_bytes(b"fake onnx"),
+    )
+    monkeypatch.setattr(tensorrt_main.TRTInferSession, "_load_engine", fail_load_engine)
+    monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", FakeEngineBuilder)
+
+    tensorrt_main.TRTInferSession(cfg)
+
+    assert build_count == 1
+
+
+def test_tensorrt_falls_back_to_rebuild_when_cached_engine_load_fails(
+    monkeypatch, tmp_path
+):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    cfg = make_tensorrt_cfg(tmp_path, engine_cfg={"cache_dir": tmp_path / "trt_cache"})
+    engine_path = tmp_path / "trt_cache" / "en_PP-OCRv4_det_mobile_sm87_fp16.engine"
+    engine_path.parent.mkdir()
+    engine_path.write_bytes(b"bad cached engine")
+
+    build_count = 0
+
+    class FakeEngineBuilder:
+        def __init__(self, **kwargs):
+            pass
+
+        def build(self):
+            nonlocal build_count
+            build_count += 1
+            return FakeEngine()
+
+    def fake_load_engine(self, path):
+        raise RuntimeError("cannot deserialize")
+
+    monkeypatch.setattr(
+        tensorrt_main.DownloadFile,
+        "run",
+        lambda input_params: Path(input_params.save_path).write_bytes(b"fake onnx"),
+    )
+    monkeypatch.setattr(tensorrt_main.TRTInferSession, "_load_engine", fake_load_engine)
+    monkeypatch.setattr(tensorrt_main, "TRTEngineBuilder", FakeEngineBuilder)
+
+    tensorrt_main.TRTInferSession(cfg)
+
+    assert build_count == 1
+
+
+def test_tensorrt_square_padding_and_crop(monkeypatch):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    session = object.__new__(tensorrt_main.TRTInferSession)
+    session._closed = True
+    session._max_square_size = 64
+
+    import numpy as np
+
+    img = np.ones((1, 3, 32, 48), dtype=np.float32)
+    padded, original_hw = session._pad_to_square(img)
+
+    assert padded.shape == (1, 3, 64, 64)
+    assert original_hw == (32, 48)
+    np.testing.assert_array_equal(padded[:, :, :32, :48], img)
+    assert padded[:, :, 32:, :].sum() == 0
+    assert padded[:, :, :, 48:].sum() == 0
+
+    cropped = session._crop_output(padded, original_hw)
+    assert cropped.shape == (1, 3, 32, 48)
+
+
+def test_tensorrt_get_dict_key_url_uses_fallback_engines(monkeypatch):
+    tensorrt_main = import_tensorrt_main(monkeypatch)
+
+    from rapidocr.inference_engine.base import FileInfo
+    from rapidocr.utils.typings import LangRec
+
+    file_info = FileInfo(
+        engine_type=EngineType.TENSORRT,
+        ocr_version=OCRVersion.PPOCRV5,
+        task_type=TaskType.REC,
+        lang_type=LangRec.EN,
+        model_type=ModelType.MOBILE,
+    )
+
+    dict_url = tensorrt_main.TRTInferSession.get_dict_key_url(file_info)
+
+    assert dict_url.endswith("/ppocrv5_en_dict.txt")