chore: update files

python/rapidocr/inference_engine/base.py

@@ -161,13 +161,14 @@ class InferSession(abc.ABC):
             f"  ocr_version   = {ocr_version}\n"
             f"  task_type     = {task_type}\n"
             f"  lang_type     = {lang_type}\n"
+            f"  model_type     = {model_type}\n"
             "\n"
             "Please refer to the official model list for supported combinations:\n"
             "https://rapidai.github.io/RapidOCRDocs/main/model_list/\n"
             "\n"
             "Example valid usage:\n"
             "  from rapidocr import LangRec, OCRVersion, RapidOCR\n"
-            "  engine = RapidOCR(params={'Rec.ocr_version': OCRVersion.PPOCRV5, 'Rec.lang_type': LangRec.CH})",
+            "  engine = RapidOCR(params={'Rec.ocr_version': OCRVersion.PPOCRV5, 'Rec.lang_type': LangRec.CH, 'Rec.model_type': 'mobile'})",
         )
         raise ValueError("Invalid OCR configuration.")
 

python/rapidocr/inference_engine/mnn/main.py

@@ -11,6 +11,7 @@ from omegaconf import DictConfig
 
 from ...utils.download_file import DownloadFile, DownloadFileInput
 from ...utils.log import logger
+from ...utils.utils import mkdir
 from ..base import FileInfo, InferSession
 
 
@@ -25,10 +26,7 @@ class MNNInferSession(InferSession):
                     "Either model_path or model_root_dir must be provided in the configuration."
                 )
             model_root_dir = Path(model_root_dir)
-            if not model_root_dir.exists():
-                raise FileNotFoundError(
-                    f"model_root_dir {model_root_dir} does not exist"
-                )
+            mkdir(model_root_dir)
 
             model_info = self.get_model_url(
                 FileInfo(

python/rapidocr/inference_engine/onnxruntime/main.py

@@ -11,6 +11,7 @@ from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions
 
 from ...utils.download_file import DownloadFile, DownloadFileInput
 from ...utils.log import logger
+from ...utils.utils import mkdir
 from ..base import FileInfo, InferSession
 from .provider_config import ProviderConfig
 
@@ -38,10 +39,7 @@ class OrtInferSession(InferSession):
                     "Either model_path or model_root_dir must be provided in the configuration."
                 )
             model_root_dir = Path(model_root_dir)
-            if not model_root_dir.exists():
-                raise FileNotFoundError(
-                    f"model_root_dir {model_root_dir} does not exist"
-                )
+            mkdir(model_root_dir)
 
             # 说明用户没有指定自己模型，使用默认模型
             model_info = self.get_model_url(

python/rapidocr/inference_engine/openvino/main.py

@@ -15,6 +15,7 @@ except ImportError:
 
 from ...utils.download_file import DownloadFile, DownloadFileInput
 from ...utils.log import logger
+from ...utils.utils import mkdir
 from ..base import FileInfo, InferSession
 from .device_config import CPUConfig
 
@@ -32,10 +33,7 @@ class OpenVINOInferSession(InferSession):
                     "Either model_path or model_root_dir must be provided in the configuration."
                 )
             model_root_dir = Path(model_root_dir)
-            if not model_root_dir.exists():
-                raise FileNotFoundError(
-                    f"model_root_dir {model_root_dir} does not exist"
-                )
+            mkdir(model_root_dir)
 
             model_info = self.get_model_url(
                 FileInfo(

python/rapidocr/inference_engine/paddle/main.py

@@ -11,6 +11,7 @@ from paddle import inference
 from ...utils.download_file import DownloadFile, DownloadFileInput
 from ...utils.log import logger
 from ...utils.typings import OCRVersion
+from ...utils.utils import mkdir
 from ..base import FileInfo, InferSession
 from .device_config import DeviceConfig
 
@@ -18,10 +19,7 @@ from .device_config import DeviceConfig
 class PaddleInferSession(InferSession):
     def __init__(self, cfg, mode: Optional[str] = None) -> 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"
-            )
+        mkdir(self.model_root_dir)
 
         self.mode = mode
 

python/rapidocr/inference_engine/pytorch/main.py

@@ -7,6 +7,7 @@ from typing import List
 import numpy as np
 import torch
 
+from ...utils.utils import mkdir
 from ..base import InferSession
 from .device_config import DeviceConfig
 from .networks.main import ModelLoader
@@ -15,10 +16,7 @@ from .networks.main import ModelLoader
 class TorchInferSession(InferSession):
     def __init__(self, cfg) -> 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"
-            )
+        mkdir(self.model_root_dir)
 
         self.device = DeviceConfig(cfg).setup_device()
         self.predictor = ModelLoader(cfg, self.device).predictor

python/rapidocr/inference_engine/tensorrt/main.py

@@ -13,6 +13,7 @@ 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
@@ -291,10 +292,7 @@ class TRTInferSession(InferSession):
                         "Either model_path or model_root_dir must be provided in the configuration."
                     )
                 model_root_dir = Path(model_root_dir)
-                if not model_root_dir.exists():
-                    raise FileNotFoundError(
-                        f"model_root_dir {model_root_dir} does not exist"
-                    )
+                mkdir(model_root_dir)
 
                 self.model_root_dir = model_root_dir
                 if not self.model_root_dir.exists():

python/tests/test_default_config.py

@@ -0,0 +1,274 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from pathlib import Path
+
+import pytest
+
+from rapidocr import RapidOCR
+from rapidocr.inference_engine.base import FileInfo, InferSession
+from rapidocr.utils.parse_parameters import ParseParams
+from rapidocr.utils.typings import (
+    EngineType,
+    LangCls,
+    LangDet,
+    LangRec,
+    ModelType,
+    OCRVersion,
+    TaskType,
+)
+
+root_dir = Path(__file__).resolve().parent.parent
+config_path = root_dir / "rapidocr" / "config.yaml"
+tests_dir = root_dir / "tests" / "test_files"
+
+
+@pytest.fixture()
+def default_cfg():
+    return ParseParams.load(config_path)
+
+
+@pytest.fixture(scope="module")
+def default_engine():
+    return RapidOCR()
+
+
+def get_default_model_info(model_cfg):
+    return InferSession.get_model_url(
+        FileInfo(
+            model_cfg.engine_type,
+            model_cfg.ocr_version,
+            model_cfg.task_type,
+            model_cfg.lang_type,
+            model_cfg.model_type,
+        )
+    )
+
+
+def test_default_global_config(default_cfg):
+    assert default_cfg.Global.text_score == 0.5
+    assert default_cfg.Global.use_det is True
+    assert default_cfg.Global.use_cls is True
+    assert default_cfg.Global.use_rec is True
+    assert default_cfg.Global.min_height == 30
+    assert default_cfg.Global.width_height_ratio == 8
+    assert default_cfg.Global.max_side_len == 2000
+    assert default_cfg.Global.min_side_len == 30
+    assert default_cfg.Global.return_word_box is False
+    assert default_cfg.Global.return_single_char_box is False
+    assert default_cfg.Global.font_path is None
+    assert default_cfg.Global.log_level == "info"
+    assert default_cfg.Global.model_root_dir is None
+
+
+def test_default_onnxruntime_engine_config(default_cfg):
+    cfg = default_cfg.EngineConfig.onnxruntime
+
+    assert cfg.intra_op_num_threads == -1
+    assert cfg.inter_op_num_threads == -1
+    assert cfg.enable_cpu_mem_arena is False
+    assert cfg.cpu_ep_cfg.arena_extend_strategy == "kSameAsRequested"
+    assert cfg.use_cuda is False
+    assert cfg.cuda_ep_cfg.device_id == 0
+    assert cfg.cuda_ep_cfg.arena_extend_strategy == "kNextPowerOfTwo"
+    assert cfg.cuda_ep_cfg.cudnn_conv_algo_search == "EXHAUSTIVE"
+    assert cfg.cuda_ep_cfg.do_copy_in_default_stream is True
+    assert cfg.use_dml is False
+    assert cfg.dml_ep_cfg is None
+    assert cfg.use_cann is False
+    assert cfg.cann_ep_cfg.device_id == 0
+    assert cfg.cann_ep_cfg.npu_mem_limit == 21474836480
+    assert cfg.cann_ep_cfg.op_select_impl_mode == "high_performance"
+    assert cfg.cann_ep_cfg.optypelist_for_implmode == "Gelu"
+    assert cfg.cann_ep_cfg.enable_cann_graph is True
+    assert cfg.use_coreml is False
+    assert cfg.coreml_ep_cfg.ModelFormat == "MLProgram"
+    assert cfg.coreml_ep_cfg.MLComputeUnits == "ALL"
+    assert cfg.coreml_ep_cfg.ModelCacheDirectory == "/tmp/RapidOCR"
+
+
+def test_default_openvino_engine_config(default_cfg):
+    cfg = default_cfg.EngineConfig.openvino
+
+    assert cfg.inference_num_threads == -1
+    assert cfg.performance_hint is None
+    assert cfg.performance_num_requests == -1
+    assert cfg.enable_cpu_pinning is None
+    assert cfg.num_streams == -1
+    assert cfg.enable_hyper_threading is None
+    assert cfg.scheduling_core_type is None
+
+
+def test_default_other_engine_configs(default_cfg):
+    assert default_cfg.EngineConfig.paddle.cpu_math_library_num_threads == -1
+    assert default_cfg.EngineConfig.paddle.use_npu is False
+    assert default_cfg.EngineConfig.paddle.npu_ep_cfg.device_id == 0
+    assert default_cfg.EngineConfig.paddle.use_cuda is False
+    assert default_cfg.EngineConfig.paddle.cuda_ep_cfg.device_id == 0
+    assert default_cfg.EngineConfig.paddle.cuda_ep_cfg.gpu_mem == 500
+
+    assert default_cfg.EngineConfig.torch.use_cuda is False
+    assert default_cfg.EngineConfig.torch.cuda_ep_cfg.device_id == 0
+    assert default_cfg.EngineConfig.torch.use_npu is False
+    assert default_cfg.EngineConfig.torch.npu_ep_cfg.device_id == 0
+    assert default_cfg.EngineConfig.torch.use_mps is False
+
+    assert default_cfg.EngineConfig.tensorrt.device_id == 0
+    assert default_cfg.EngineConfig.tensorrt.use_fp16 is True
+    assert default_cfg.EngineConfig.tensorrt.use_int8 is False
+    assert default_cfg.EngineConfig.tensorrt.workspace_size == 1073741824
+    assert default_cfg.EngineConfig.tensorrt.cache_dir is None
+    assert default_cfg.EngineConfig.tensorrt.force_rebuild is False
+    assert default_cfg.EngineConfig.tensorrt.det_profile.opt_shape == [
+        1,
+        3,
+        736,
+        736,
+    ]
+    assert default_cfg.EngineConfig.tensorrt.rec_profile.opt_shape == [
+        6,
+        3,
+        48,
+        320,
+    ]
+    assert default_cfg.EngineConfig.tensorrt.cls_profile.opt_shape == [
+        6,
+        3,
+        48,
+        192,
+    ]
+
+    assert default_cfg.EngineConfig.mnn == {}
+
+
+@pytest.mark.parametrize(
+    "section,task_type,expected_model_name,expected_sha256",
+    [
+        (
+            "Det",
+            TaskType.DET,
+            "PP-OCRv6_det_small.onnx",
+            "090f04abcd9d9a7498bc4ebf677e4cb9bdce1fe4197ddb7e529f1ef44e1ff94f",
+        ),
+        (
+            "Rec",
+            TaskType.REC,
+            "PP-OCRv6_rec_small.onnx",
+            "6f327246b50388f3c176ae304bd95767ea6dc0c9ae92153ef8cbe210b3c14884",
+        ),
+    ],
+)
+def test_default_ppocrv6_model_config(
+    default_cfg, section, task_type, expected_model_name, expected_sha256
+):
+    model_cfg = default_cfg[section]
+
+    assert model_cfg.engine_type == EngineType.ONNXRUNTIME
+    assert model_cfg.ocr_version == OCRVersion.PPOCRV6
+    assert model_cfg.model_type == ModelType.SMALL
+    assert model_cfg.task_type == task_type
+
+    model_info = get_default_model_info(model_cfg)
+
+    assert model_info["model_dir"].endswith(
+        f"/PP-OCRv6/{task_type.value}/{expected_model_name}"
+    )
+    assert model_info["SHA256"] == expected_sha256
+
+
+def test_default_det_config(default_cfg):
+    assert default_cfg.Det.engine_type == EngineType.ONNXRUNTIME
+    assert default_cfg.Det.lang_type == LangDet.CH
+    assert default_cfg.Det.model_type == ModelType.SMALL
+    assert default_cfg.Det.ocr_version == OCRVersion.PPOCRV6
+    assert default_cfg.Det.task_type == TaskType.DET
+    assert default_cfg.Det.model_path is None
+    assert default_cfg.Det.model_dir is None
+    assert default_cfg.Det.limit_side_len == 736
+    assert default_cfg.Det.limit_type == "min"
+    assert default_cfg.Det.std == [0.5, 0.5, 0.5]
+    assert default_cfg.Det.mean == [0.5, 0.5, 0.5]
+    assert default_cfg.Det.thresh == 0.3
+    assert default_cfg.Det.box_thresh == 0.5
+    assert default_cfg.Det.max_candidates == 1000
+    assert default_cfg.Det.unclip_ratio == 1.6
+    assert default_cfg.Det.use_dilation is True
+    assert default_cfg.Det.score_mode == "fast"
+
+
+def test_default_cls_config(default_cfg):
+    assert default_cfg.Cls.engine_type == EngineType.ONNXRUNTIME
+    assert default_cfg.Cls.lang_type == LangCls.CH
+    assert default_cfg.Cls.model_type == ModelType.MOBILE
+    assert default_cfg.Cls.ocr_version == OCRVersion.PPOCRV4
+    assert default_cfg.Cls.task_type == TaskType.CLS
+    assert default_cfg.Cls.model_path is None
+    assert default_cfg.Cls.model_dir is None
+    assert default_cfg.Cls.cls_image_shape == [3, 48, 192]
+    assert default_cfg.Cls.cls_batch_num == 6
+    assert default_cfg.Cls.cls_thresh == 0.9
+    assert default_cfg.Cls.label_list == ["0", "180"]
+
+
+def test_default_rec_config(default_cfg):
+    assert default_cfg.Rec.engine_type == EngineType.ONNXRUNTIME
+    assert default_cfg.Rec.lang_type == LangRec.CH
+    assert default_cfg.Rec.model_type == ModelType.SMALL
+    assert default_cfg.Rec.ocr_version == OCRVersion.PPOCRV6
+    assert default_cfg.Rec.task_type == TaskType.REC
+    assert default_cfg.Rec.model_path is None
+    assert default_cfg.Rec.model_dir is None
+    assert default_cfg.Rec.rec_keys_path is None
+    assert default_cfg.Rec.rec_img_shape == [3, 48, 320]
+    assert default_cfg.Rec.rec_batch_num == 6
+
+
+def test_default_cls_model_config(default_cfg):
+    model_info = get_default_model_info(default_cfg.Cls)
+
+    assert model_info["model_dir"].endswith(
+        "/PP-OCRv4/cls/ch_ppocr_mobile_v2.0_cls_mobile.onnx"
+    )
+    assert (
+        model_info["SHA256"]
+        == "e47acedf663230f8863ff1ab0e64dd2d82b838fceb5957146dab185a89d6215c"
+    )
+
+
+def test_default_engine_runtime_config(default_engine):
+    assert default_engine.use_det is True
+    assert default_engine.use_cls is True
+    assert default_engine.use_rec is True
+    assert default_engine.cfg.Det.ocr_version == OCRVersion.PPOCRV6
+    assert default_engine.cfg.Det.model_type == ModelType.SMALL
+    assert default_engine.cfg.Rec.ocr_version == OCRVersion.PPOCRV6
+    assert default_engine.cfg.Rec.model_type == ModelType.SMALL
+    assert default_engine.cfg.Cls.ocr_version == OCRVersion.PPOCRV4
+    assert default_engine.cfg.Cls.model_type == ModelType.MOBILE
+    assert default_engine.text_det.limit_side_len == 736
+    assert default_engine.text_det.limit_type == "min"
+    assert default_engine.text_cls.cls_image_shape == [3, 48, 192]
+    assert default_engine.text_rec.rec_image_shape == [3, 48, 320]
+    assert default_engine.text_rec.rec_batch_num == 6
+
+
+def test_default_engine_real_full_pipeline(default_engine):
+    result = default_engine(tests_dir / "ch_en_num.jpg")
+
+    assert result.txts is not None
+    assert len(result) == 18
+    assert result.txts[:3] == ("正品促销", "大桶装更划算", "强力去污符合国标")
+
+
+def test_default_engine_real_rec_only(default_engine):
+    result = default_engine(
+        tests_dir / "text_rec.jpg",
+        use_det=False,
+        use_cls=False,
+        use_rec=True,
+    )
+
+    assert result.txts is not None
+    assert len(result) == 1
+    assert result.txts[0] == "韩国小馆"
+    assert result.scores[0] > 0.99

python/tests/test_prepare_wheel_assets.py

@@ -0,0 +1,70 @@
+# -*- 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 tools.prepare_wheel_assets import (
+    ModelSpec,
+    load_yaml,
+    resolve_assets,
+    resolve_model_specs,
+    select_model_info,
+)
+
+config_path = root_dir / "rapidocr" / "config.yaml"
+models_path = root_dir / "rapidocr" / "default_models.yaml"
+
+
+@pytest.fixture(scope="module")
+def default_config():
+    return load_yaml(config_path)
+
+
+@pytest.fixture(scope="module")
+def model_registry():
+    return load_yaml(models_path)
+
+
+def test_default_wheel_assets_use_routed_ppocrv6_models(
+    default_config, model_registry
+):
+    specs = resolve_model_specs(default_config)
+    assets = resolve_assets(specs, model_registry)
+    asset_paths = {asset.relative_path.as_posix() for asset in assets}
+
+    assert asset_paths == {
+        "PP-OCRv6_det_small.onnx",
+        "ch_ppocr_mobile_v2.0_cls_mobile.onnx",
+        "PP-OCRv6_rec_small.onnx",
+    }
+
+
+@pytest.mark.parametrize(
+    "spec,expected_model_name",
+    [
+        (
+            ModelSpec("onnxruntime", "PP-OCRv6", "det", "ch", "small"),
+            "PP-OCRv6_det_small.onnx",
+        ),
+        (
+            ModelSpec("onnxruntime", "PP-OCRv6", "rec", "ch", "small"),
+            "PP-OCRv6_rec_small.onnx",
+        ),
+        (
+            ModelSpec("onnxruntime", "PP-OCRv4", "cls", "ch", "mobile"),
+            "ch_ppocr_mobile_v2.0_cls_mobile.onnx",
+        ),
+    ],
+)
+def test_select_model_info_supports_default_model_specs(
+    model_registry, spec, expected_model_name
+):
+    model_info = select_model_info(spec, model_registry)
+
+    assert model_info["model_dir"].endswith(expected_model_name)

python/tests/test_rec_language.py

@@ -10,13 +10,7 @@ from pytest import mark
 root_dir = Path(__file__).resolve().parent.parent
 sys.path.append(str(root_dir))
 
-from rapidocr import (
-    EngineType,
-    LangRec,
-    ModelType,
-    OCRVersion,
-    RapidOCR,
-)
+from rapidocr import EngineType, LangRec, ModelType, OCRVersion, RapidOCR
 
 tests_dir = root_dir / "tests" / "test_files"
 img_path = tests_dir / "ch_en_num.jpg"
@@ -109,9 +103,7 @@ def test_th_lang(engine_type):
     img_path = tests_dir / "th_rec.jpg"
     result = engine(img_path, use_det=False, use_cls=False, use_rec=True)
     assert result.txts is not None
-    assert (
-        result.txts[0] == "การแพทย์แผนจีนช่วยให้เด็กสมองพิการในคาซัคสถานกลับมาเดินได้"
-    )
+    assert result.txts[0] == "การแพทย์แผนจีนช่วยให้เด็กสมองพิการในคาซัคสถานกลับมาเดินได้"
 
 
 @mark.parametrize(

python/tools/prepare_wheel_assets.py

@@ -10,6 +10,11 @@ from pathlib import Path
 from typing import Iterable, List, Optional
 
 PYTHON_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PYTHON_ROOT))
+
+from rapidocr.utils.model_resolver import normalize_lang, resolve_model_key
+from rapidocr.utils.typings import ModelType, OCRVersion, TaskType
+
 DEFAULT_CONFIG_YAML = PYTHON_ROOT / "rapidocr" / "config.yaml"
 DEFAULT_MANIFEST_IN = PYTHON_ROOT / "MANIFEST.in"
 DEFAULT_MODEL_DIR = PYTHON_ROOT / "rapidocr" / "models"
@@ -147,8 +152,28 @@ def select_model_info(spec: ModelSpec, registry: dict) -> dict:
             f"{spec.engine}.{spec.ocr_version}.{spec.task}"
         ) from e
 
+    model_key = _resolve_model_key(spec)
+    if model_key is not None:
+        model_info = task_models.get(model_key)
+        if not isinstance(model_info, dict):
+            raise RuntimeError(
+                "No model registry entry matches routed default config: "
+                f"{spec.engine}.{spec.ocr_version}.{spec.task} "
+                f"lang={spec.lang} model_type={spec.model_type} "
+                f"model_key={model_key}"
+            )
+        return model_info
+
+    lang = normalize_lang(spec.lang)
+    if spec.model_type == ModelType.SERVER.value:
+        for model_name, model_info in task_models.items():
+            if model_name.startswith(lang) and spec.model_type in model_name:
+                if not isinstance(model_info, dict):
+                    raise RuntimeError(f"Invalid model registry entry: {model_name}")
+                return model_info
+
     for model_name, model_info in task_models.items():
-        if model_name.startswith(spec.lang) and spec.model_type in model_name:
+        if model_name.startswith(lang) and spec.model_type in model_name:
             if not isinstance(model_info, dict):
                 raise RuntimeError(f"Invalid model registry entry: {model_name}")
             return model_info
@@ -160,6 +185,17 @@ def select_model_info(spec: ModelSpec, registry: dict) -> dict:
     )
 
 
+def _resolve_model_key(spec: ModelSpec) -> Optional[str]:
+    try:
+        task_type = TaskType(spec.task)
+        ocr_version = OCRVersion(spec.ocr_version)
+        model_type = ModelType(spec.model_type)
+    except ValueError:
+        return None
+
+    return resolve_model_key(task_type, ocr_version, spec.lang, model_type)
+
+
 def model_info_to_assets(spec: ModelSpec, model_info: dict) -> List[WheelAsset]:
     model_url = _required_str(model_info, "model_dir", "model registry")