diff --git a/python/rapidocr/__init__.py b/python/rapidocr/__init__.py index 1b47892..50ad245 100644 --- a/python/rapidocr/__init__.py +++ b/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", ] diff --git a/python/rapidocr/config.yaml b/python/rapidocr/config.yaml index 18850a1..7327e79 100644 --- a/python/rapidocr/config.yaml +++ b/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 diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index 6edaa89..128dfc1 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/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): diff --git a/python/rapidocr/inference_engine/mnn/main.py b/python/rapidocr/inference_engine/mnn/main.py index baef27d..e776370 100644 --- a/python/rapidocr/inference_engine/mnn/main.py +++ b/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"], diff --git a/python/rapidocr/inference_engine/onnxruntime/main.py b/python/rapidocr/inference_engine/onnxruntime/main.py index 7ed03b7..e678338 100644 --- a/python/rapidocr/inference_engine/onnxruntime/main.py +++ b/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"], diff --git a/python/rapidocr/inference_engine/openvino/main.py b/python/rapidocr/inference_engine/openvino/main.py index ce74abc..280d5b9 100644 --- a/python/rapidocr/inference_engine/openvino/main.py +++ b/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"], diff --git a/python/rapidocr/inference_engine/paddle/main.py b/python/rapidocr/inference_engine/paddle/main.py index e6f8b71..cb82981 100644 --- a/python/rapidocr/inference_engine/paddle/main.py +++ b/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( diff --git a/python/rapidocr/inference_engine/pytorch/main.py b/python/rapidocr/inference_engine/pytorch/main.py index 40f9535..561a34e 100644 --- a/python/rapidocr/inference_engine/pytorch/main.py +++ b/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 diff --git a/python/rapidocr/inference_engine/pytorch/networks/main.py b/python/rapidocr/inference_engine/pytorch/networks/main.py index c40bdcf..8d0b050 100644 --- a/python/rapidocr/inference_engine/pytorch/networks/main.py +++ b/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, diff --git a/python/rapidocr/inference_engine/tensorrt/main.py b/python/rapidocr/inference_engine/tensorrt/main.py index 121264e..083b612 100644 --- a/python/rapidocr/inference_engine/tensorrt/main.py +++ b/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) diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 3614a56..9cfa273 100644 --- a/python/rapidocr/main.py +++ b/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, diff --git a/python/rapidocr/utils/download_models.py b/python/rapidocr/utils/download_models.py new file mode 100644 index 0000000..438a88e --- /dev/null +++ b/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, + ) + ) diff --git a/python/tests/base_module.py b/python/tests/base_module.py deleted file mode 100644 index e538494..0000000 --- a/python/tests/base_module.py +++ /dev/null @@ -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" diff --git a/python/tests/test_cli.py b/python/tests/test_cli.py index 84b4048..e5ee7ab 100644 --- a/python/tests/test_cli.py +++ b/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)) diff --git a/python/tests/test_files/test_config.yaml b/python/tests/test_files/test_config.yaml new file mode 100644 index 0000000..6f6ccec --- /dev/null +++ b/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