feat(python): add model_root_dir and download_models (API + CLI)

python/rapidocr/__init__.py

@@ -2,6 +2,7 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 from .main import RapidOCR
+from .utils.download_models import download_models
 from .utils.load_image import LoadImageError
 from .utils.typings import EngineType, LangCls, LangDet, LangRec, ModelType, OCRVersion
 from .utils.vis_res import VisRes
@@ -16,4 +17,5 @@ __all__ = [
     "ModelType",
     "OCRVersion",
     "VisRes",
+    "download_models",
 ]

python/rapidocr/config.yaml

@@ -16,6 +16,8 @@ Global:
     font_path: null
     log_level: "info" # debug / info / warning / error / critical
 
+    model_root_dir: null
+
 EngineConfig:
     onnxruntime:
         intra_op_num_threads: -1
@@ -45,7 +47,7 @@ EngineConfig:
             enable_cann_graph: true
 
         use_coreml: false
-        coreml_ep_cfg: 
+        coreml_ep_cfg:
             ModelFormat: "MLProgram"
             MLComputeUnits: "ALL"
             RequireStaticInputShapes: 0

python/rapidocr/inference_engine/base.py

@@ -83,7 +83,6 @@ class FileInfo:
 
 class InferSession(abc.ABC):
     model_info = OmegaConf.load(MODEL_URL_PATH)
-    DEFAULT_MODEL_PATH = cur_dir / "models"
 
     @abc.abstractmethod
     def __init__(self, config):

python/rapidocr/inference_engine/mnn/main.py

@@ -16,7 +16,9 @@ from ..base import FileInfo, InferSession
 
 class MNNInferSession(InferSession):
     def __init__(self, cfg: DictConfig):
-        super().__init__(cfg)
+        model_root_dir = Path(cfg.get("model_root_dir"))
+        if not model_root_dir.exists():
+            raise FileNotFoundError(f"model_root_dir {model_root_dir} does not exist")
 
         model_path = cfg.get("model_path", None)
         if model_path is None:
@@ -29,7 +31,7 @@ class MNNInferSession(InferSession):
                     model_type=cfg.model_type,
                 )
             )
-            model_path = self.DEFAULT_MODEL_PATH / Path(model_info["model_dir"]).name
+            model_path = model_root_dir / Path(model_info["model_dir"]).name
             DownloadFile.run(
                 DownloadFileInput(
                     file_url=model_info["model_dir"],

python/rapidocr/inference_engine/onnxruntime/main.py

@@ -17,6 +17,10 @@ from .provider_config import ProviderConfig
 
 class OrtInferSession(InferSession):
     def __init__(self, cfg: Dict[str, Any]):
+        model_root_dir = Path(cfg.get("model_root_dir"))
+        if not model_root_dir.exists():
+            raise FileNotFoundError(f"model_root_dir {model_root_dir} does not exist")
+
         # support custom session (PR #451)
         session = cfg.get("session", None)
         if session is not None:
@@ -41,7 +45,7 @@ class OrtInferSession(InferSession):
                     model_type=cfg.model_type,
                 )
             )
-            model_path = self.DEFAULT_MODEL_PATH / Path(model_info["model_dir"]).name
+            model_path = model_root_dir / Path(model_info["model_dir"]).name
             download_params = DownloadFileInput(
                 file_url=model_info["model_dir"],
                 sha256=model_info["SHA256"],

python/rapidocr/inference_engine/openvino/main.py

@@ -21,7 +21,9 @@ from .device_config import CPUConfig
 
 class OpenVINOInferSession(InferSession):
     def __init__(self, cfg: DictConfig):
-        super().__init__(cfg)
+        model_root_dir = Path(cfg.get("model_root_dir"))
+        if not model_root_dir.exists():
+            raise FileNotFoundError(f"model_root_dir {model_root_dir} does not exist")
 
         core = Core()
 
@@ -36,7 +38,7 @@ class OpenVINOInferSession(InferSession):
                     model_type=cfg.model_type,
                 )
             )
-            model_path = self.DEFAULT_MODEL_PATH / Path(model_info["model_dir"]).name
+            model_path = model_root_dir / Path(model_info["model_dir"]).name
             download_params = DownloadFileInput(
                 file_url=model_info["model_dir"],
                 sha256=model_info["SHA256"],

python/rapidocr/inference_engine/paddle/main.py

@@ -17,6 +17,12 @@ 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"
+            )
+
         self.mode = mode
 
         pdmodel_path, pdiparams_path = self.setup_model(cfg)
@@ -113,7 +119,7 @@ class PaddleInferSession(InferSession):
     ) -> Path:
         model_file_url = f"{default_model_dir}/{model_file_name}"
         model_file_path = (
-            self.DEFAULT_MODEL_PATH / Path(default_model_dir).name / model_file_name
+            self.model_root_dir / Path(default_model_dir).name / model_file_name
         )
         DownloadFile.run(
             DownloadFileInput(

python/rapidocr/inference_engine/pytorch/main.py

@@ -13,6 +13,12 @@ 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"
+            )
+
         self.device = DeviceConfig(cfg).setup_device()
         self.predictor = ModelLoader(cfg, self.device).predictor
 
@@ -21,7 +27,7 @@ class TorchInferSession(InferSession):
             inp = torch.from_numpy(img)
             inp = inp.to(self.device)
             outputs = self.predictor(inp).cpu().numpy()
-            if self.device.type=="mps":
+            if self.device.type == "mps":
                 del inp
                 torch.mps.empty_cache()
             return outputs

python/rapidocr/inference_engine/pytorch/networks/main.py

@@ -37,7 +37,7 @@ class ModelLoader:
                 )
             )
             default_model_url = model_info["model_dir"]
-            model_path = InferSession.DEFAULT_MODEL_PATH / Path(default_model_url).name
+            model_path = cfg.model_root_dir / Path(default_model_url).name
             DownloadFile.run(
                 DownloadFileInput(
                     file_url=default_model_url,

python/rapidocr/inference_engine/tensorrt/main.py

@@ -74,6 +74,12 @@ class TRTInferSession(InferSession):
             AssertionError: If CUDA device setup fails.
             RuntimeError: If engine building fails.
         """
+        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"
+            )
+
         self.cfg = cfg
         self.engine_cfg = cfg.get("engine_cfg", {})
         self._closed = False
@@ -456,7 +462,7 @@ class TRTInferSession(InferSession):
         """
         cache_dir = self.engine_cfg.get("cache_dir")
         if cache_dir is None:
-            cache_dir = self.DEFAULT_MODEL_PATH
+            cache_dir = self.model_root_dir / "models"
 
         cache_dir = Path(cache_dir)
         cache_dir.mkdir(parents=True, exist_ok=True)

python/rapidocr/main.py

@@ -15,6 +15,7 @@ from .ch_ppocr_cls import TextClassifier, TextClsOutput
 from .ch_ppocr_det import TextDetector, TextDetOutput
 from .ch_ppocr_rec import TextRecInput, TextRecognizer, TextRecOutput
 from .cli import check_install, generate_cfg
+from .utils.download_models import download_models
 from .utils.load_image import LoadImage
 from .utils.log import logger
 from .utils.output import RapidOCROutput
@@ -54,6 +55,9 @@ class RapidOCR:
 
         if params:
             cfg = ParseParams.update_batch(cfg, params)
+
+        if cfg.Global.model_root_dir is None:
+            cfg.Global.model_root_dir = root_dir / "models"
         return cfg
 
     def _initialize(self, cfg: DictConfig):
@@ -63,15 +67,18 @@ class RapidOCR:
 
         self.use_det = cfg.Global.use_det
         cfg.Det.engine_cfg = cfg.EngineConfig[cfg.Det.engine_type.value]
+        cfg.Det.model_root_dir = cfg.Global.model_root_dir
         self.text_det = TextDetector(cfg.Det)
 
         self.use_cls = cfg.Global.use_cls
         cfg.Cls.engine_cfg = cfg.EngineConfig[cfg.Cls.engine_type.value]
+        cfg.Cls.model_root_dir = cfg.Global.model_root_dir
         self.text_cls = TextClassifier(cfg.Cls)
 
         self.use_rec = cfg.Global.use_rec
         cfg.Rec.engine_cfg = cfg.EngineConfig[cfg.Rec.engine_type.value]
         cfg.Rec.font_path = cfg.Global.font_path
+        cfg.Rec.model_root_dir = cfg.Global.model_root_dir
         self.text_rec = TextRecognizer(cfg.Rec)
 
         self.load_img = LoadImage()
@@ -372,6 +379,10 @@ def parse_args(arg_list: Optional[List[str]] = None):
     )
     parser_check.set_defaults(func=check_install)
 
+    parser_download = subparser.add_parser("download_models", help="Download models")
+    parser_download.add_argument("--config", type=str)
+    parser_download.set_defaults(func=download_models)
+
     args = parser.parse_args(arg_list)
     return args
 
@@ -383,6 +394,10 @@ def main(arg_list: Optional[List[str]] = None):
         generate_cfg(args)
         return
 
+    if args.command == "download_models":
+        download_models(args.config)
+        return
+
     params = {
         "Global.text_score": args.text_score,
         "Global.return_word_box": args.return_word_box,

python/rapidocr/utils/download_models.py

@@ -0,0 +1,126 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from pathlib import Path
+from typing import Union
+
+from ..inference_engine.base import FileInfo, InferSession
+from .download_file import DownloadFile, DownloadFileInput
+from .log import logger
+from .parse_parameters import ParseParams
+from .typings import EngineType, TaskType
+from .utils import mkdir
+
+root_dir = Path(__file__).resolve().parent.parent
+
+
+def download_models(config_path: Union[str, Path]) -> None:
+    cfg = ParseParams.load(config_path)
+
+    model_root_dir = cfg.Global.get("model_root_dir")
+    if model_root_dir is None:
+        model_root_dir = root_dir / "models"
+
+    model_root_dir = Path(model_root_dir).expanduser().resolve()
+    mkdir(model_root_dir)
+
+    use_det = cfg.Global.get("use_det", True)
+    use_cls = cfg.Global.get("use_cls", True)
+    use_rec = cfg.Global.get("use_rec", True)
+    if not (use_det or use_cls or use_rec):
+        raise ValueError(
+            "Config has use_det, use_cls, use_rec all false; at least one must be True."
+        )
+
+    if use_det:
+        download_task(
+            model_root_dir,
+            FileInfo(
+                cfg.Det.engine_type,
+                cfg.Det.ocr_version,
+                cfg.Det.task_type,
+                cfg.Det.lang_type,
+                cfg.Det.model_type,
+            ),
+        )
+
+    if use_cls:
+        download_task(
+            model_root_dir,
+            FileInfo(
+                cfg.Cls.engine_type,
+                cfg.Cls.ocr_version,
+                cfg.Cls.task_type,
+                cfg.Cls.lang_type,
+                cfg.Cls.model_type,
+            ),
+        )
+
+    if use_rec:
+        download_task(
+            model_root_dir,
+            FileInfo(
+                cfg.Rec.engine_type,
+                cfg.Rec.ocr_version,
+                cfg.Rec.task_type,
+                cfg.Rec.lang_type,
+                cfg.Rec.model_type,
+            ),
+        )
+
+    print(f"Models downloaded to {model_root_dir}")
+    print(
+        f"Please initialize RapidOCR using the {config_path} configuration file to ensure the model used is consistent with the pre-downloaded one."
+    )
+
+
+def download_task(cache_dir: Path, file_info: FileInfo) -> None:
+    info = InferSession.get_model_url(file_info)
+    engine_type = file_info.engine_type
+
+    if engine_type != EngineType.PADDLE:
+        download_single_file(cache_dir, info)
+    else:
+        download_paddle_bundle(cache_dir, info)
+
+    if file_info.task_type == TaskType.REC:
+        dict_url = info.get("dict_url")
+        if dict_url:
+            DownloadFile.run(
+                DownloadFileInput(
+                    file_url=dict_url,
+                    save_path=cache_dir / Path(dict_url).name,
+                    logger=logger,
+                    sha256=None,
+                )
+            )
+
+
+def download_single_file(cache_dir: Path, info: dict) -> None:
+    model_url = info["model_dir"]
+    save_path = cache_dir / Path(model_url).name
+    DownloadFile.run(
+        DownloadFileInput(
+            file_url=model_url,
+            save_path=save_path,
+            logger=logger,
+            sha256=info.get("SHA256"),
+        )
+    )
+
+
+def download_paddle_bundle(cache_dir: Path, info: dict) -> None:
+    model_url_dir = info["model_dir"].rstrip("/")
+    save_model_dir = cache_dir / Path(model_url_dir).name
+    for name, sha in info.items():
+        if name in ("model_dir", "dict_url"):
+            continue
+
+        DownloadFile.run(
+            DownloadFileInput(
+                file_url=f"{model_url_dir}/{name}",
+                save_path=save_model_dir / name,
+                logger=logger,
+                sha256=sha,
+            )
+        )

python/tests/base_module.py

@@ -1,63 +0,0 @@
-# -*- encoding: utf-8 -*-
-# @Author: SWHL
-# @Contact: liekkaskono@163.com
-import importlib
-import sys
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Optional, Union
-
-import requests
-import yaml
-from tqdm import tqdm
-
-
-class BaseModule:
-    def __init__(self, package_name: str = "rapidocr_onnxruntime"):
-        self.package_name = package_name
-        self.root_dir = Path(__file__).resolve().parent.parent
-        self.package_dir = self.root_dir / self.package_name
-        self.tests_dir = self.root_dir / "tests"
-
-        sys.path.append(str(self.root_dir))
-        sys.path.append(str(self.package_dir))
-
-    def init_module(self, module_name: str, class_name: Optional[str] = None):
-        if class_name is None:
-            module_part = importlib.import_module(f"{self.package_name}")
-            return module_part
-        module_part = importlib.import_module(f"{self.package_name}.{module_name}")
-        return getattr(module_part, class_name)
-
-    @staticmethod
-    def read_yaml(yaml_path: str):
-        with open(yaml_path, "rb") as f:
-            data = yaml.load(f, Loader=yaml.Loader)
-        return data
-
-
-def download_file(url: str, save_path: Union[str, Path]):
-    response = requests.get(url, stream=True, timeout=60)
-    status_code = response.status_code
-
-    if status_code != 200:
-        raise DownloadModelError("Something went wrong while downloading models")
-
-    total_size_in_bytes = int(response.headers.get("content-length", 1))
-    block_size = 1024  # 1 Kibibyte
-    with tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True) as pb:
-        with open(save_path, "wb") as file:
-            for data in response.iter_content(block_size):
-                pb.update(len(data))
-                file.write(data)
-
-
-class DownloadModelError(Exception):
-    pass
-
-
-@dataclass
-class Platform:
-    mac: str = "Darwin"
-    windows: str = "Windows"
-    linux: str = "Linux"

python/tests/test_cli.py

@@ -2,6 +2,7 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 import shlex
+import shutil
 import sys
 from pathlib import Path
 
@@ -31,6 +32,17 @@ def test_cli(capsys, cmd, gt):
     assert gt in output
 
 
+@mark.parametrize("cmd", [f"download_models --config {tests_dir}/test_config.yaml"])
+def test_cli_download_models(capsys, cmd):
+    main(shlex.split(cmd))
+    output = capsys.readouterr().out.strip()
+    assert "Models downloaded to" in output
+
+    model_dir = Path("./test_models").expanduser().resolve()
+    assert model_dir.exists()
+    shutil.rmtree(model_dir)
+
+
 @mark.parametrize("cmd", [f"config --save_cfg_file {tests_dir}/config.yaml"])
 def test_cli_config(capsys, cmd):
     main(shlex.split(cmd))

python/tests/test_files/test_config.yaml

@@ -0,0 +1,179 @@
+Global:
+    text_score: 0.5
+
+    use_det: false
+    use_cls: true
+    use_rec: false
+
+    min_height: 30
+    width_height_ratio: 8
+    max_side_len: 2000
+    min_side_len: 30
+
+    return_word_box: false
+    return_single_char_box: false
+
+    font_path: null
+    log_level: "info" # debug / info / warning / error / critical
+
+    model_root_dir: ./test_models
+
+EngineConfig:
+    onnxruntime:
+        intra_op_num_threads: -1
+        inter_op_num_threads: -1
+        enable_cpu_mem_arena: false
+
+        cpu_ep_cfg:
+            arena_extend_strategy: "kSameAsRequested"
+
+        use_cuda: false
+        cuda_ep_cfg:
+            device_id: 0
+            arena_extend_strategy: "kNextPowerOfTwo"
+            cudnn_conv_algo_search: "EXHAUSTIVE"
+            do_copy_in_default_stream: true
+
+        use_dml: false
+        dm_ep_cfg: null
+
+        use_cann: false
+        cann_ep_cfg:
+            device_id: 0
+            arena_extend_strategy: "kNextPowerOfTwo"
+            npu_mem_limit:  21474836480 # 20 * 1024 * 1024 * 1024
+            op_select_impl_mode: "high_performance"
+            optypelist_for_implmode: "Gelu"
+            enable_cann_graph: true
+
+        use_coreml: false
+        coreml_ep_cfg:
+            ModelFormat: "MLProgram"
+            MLComputeUnits: "ALL"
+            RequireStaticInputShapes: 0
+            EnableOnSubgraphs: 0
+            SpecializationStrategy: "FastPrediction"
+            ProfileComputePlan: 0
+            AllowLowPrecisionAccumulationOnGPU: 0
+            ModelCacheDirectory: "/tmp/RapidOCR"
+
+    openvino:
+        inference_num_threads: -1
+        performance_hint: null
+        performance_num_requests: -1
+        enable_cpu_pinning: null
+        num_streams: -1
+        enable_hyper_threading: null
+        scheduling_core_type: null
+
+    paddle:
+        cpu_math_library_num_threads: -1
+
+        use_npu: false
+        npu_ep_cfg:
+            device_id: 0
+            envs:
+                FLAGS_npu_jit_compile: 0
+                FLAGS_use_stride_kernel: 0
+                FLAGS_allocator_strategy: "auto_growth"
+                CUSTOM_DEVICE_BLACK_LIST: "pad3d,pad3d_grad,set_value,set_value_with_tensor"
+                FLAGS_npu_scale_aclnn: "True"
+                FLAGS_npu_split_aclnn: "True"
+
+        use_cuda: false
+        cuda_ep_cfg:
+            device_id: 0
+            gpu_mem: 500
+
+    torch:
+        use_cuda: false
+        cuda_ep_cfg:
+            device_id: 0
+
+        use_npu: false
+        npu_ep_cfg:
+            device_id: 0
+
+        use_mps: false
+
+    tensorrt:
+        device_id: 0
+        use_fp16: true
+        use_int8: false
+        workspace_size: 1073741824  # 1GB = 1 << 30
+
+        # Engine caching
+        cache_dir: null  # null = use default models dir
+        force_rebuild: false
+
+        # Dynamic shape optimization profiles
+        det_profile:
+            min_shape: [1, 3, 32, 32]
+            opt_shape: [1, 3, 736, 736]
+            max_shape: [1, 3, 2048, 2048]
+
+        rec_profile:
+            min_shape: [1, 3, 48, 32]
+            opt_shape: [6, 3, 48, 320]
+            max_shape: [6, 3, 48, 2048]
+
+        cls_profile:
+            min_shape: [1, 3, 48, 32]
+            opt_shape: [6, 3, 48, 192]
+            max_shape: [6, 3, 48, 192]
+
+    mnn: {}
+
+Det:
+    engine_type: "onnxruntime"
+    lang_type: "ch"
+    model_type: "mobile"
+    ocr_version: "PP-OCRv4"
+
+    task_type: "det"
+
+    model_path: null
+    model_dir: null
+
+    limit_side_len: 736
+    limit_type: min
+    std: [ 0.5, 0.5, 0.5 ]
+    mean: [ 0.5, 0.5, 0.5 ]
+
+    thresh: 0.3
+    box_thresh: 0.5
+    max_candidates: 1000
+    unclip_ratio: 1.6
+    use_dilation: true
+    score_mode: fast
+
+Cls:
+    engine_type: "onnxruntime"
+    lang_type: "ch"
+    model_type: "mobile"
+    ocr_version: "PP-OCRv4"
+
+    task_type: "cls"
+
+    model_path: null
+    model_dir: null
+
+    cls_image_shape: [3, 48, 192]
+    cls_batch_num: 6
+    cls_thresh: 0.9
+    label_list: ["0", "180"]
+
+Rec:
+    engine_type: "onnxruntime"
+    lang_type: "ch"
+    model_type: "mobile"
+    ocr_version: "PP-OCRv4"
+
+    task_type: "rec"
+
+    model_path: null
+    model_dir: null
+
+    rec_keys_path: null
+    rec_img_shape: [3, 48, 320]
+    rec_batch_num: 6