diff --git a/python/rapidocr/inference_engine/__init__.py b/python/rapidocr/inference_engine/__init__.py index ab1f42d..0ecdd4f 100644 --- a/python/rapidocr/inference_engine/__init__.py +++ b/python/rapidocr/inference_engine/__init__.py @@ -1,6 +1,3 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -from .onnxruntime import OrtInferSession -from .paddlepaddle import PaddleInferSession -from .torch import TorchInferSession diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index c524820..ed4ca28 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/python/rapidocr/inference_engine/base.py @@ -10,6 +10,7 @@ import numpy as np from omegaconf import DictConfig, OmegaConf from ..utils.logger import Logger +from ..utils.utils import import_package cur_dir = Path(__file__).resolve().parent.parent MODEL_URL_PATH = cur_dir / "default_models.yaml" @@ -29,21 +30,33 @@ def get_engine(engine_name: str): logger.info("Using engine_name: %s", engine_name) if engine_name == Engine.ONNXRUNTIME.value: + if not import_package("onnxruntime"): + raise ImportError("onnxruntime is not installed.") + from .onnxruntime import OrtInferSession return OrtInferSession if engine_name == Engine.OPENVINO.value: + if not import_package("openvino"): + raise ImportError("openvino is not installed") + from .openvino import OpenVINOInferSession return OpenVINOInferSession if engine_name == Engine.PADDLE.value: + if not import_package("paddleocr"): + raise ImportError("paddleocr is not installed") + from .paddlepaddle import PaddleInferSession return PaddleInferSession if engine_name == Engine.TORCH.value: + if not import_package("torch"): + raise ImportError("torch is not installed") + from .torch import TorchInferSession return TorchInferSession diff --git a/python/rapidocr/utils/utils.py b/python/rapidocr/utils/utils.py index 5e44a7e..4fa4830 100644 --- a/python/rapidocr/utils/utils.py +++ b/python/rapidocr/utils/utils.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +import importlib from pathlib import Path from typing import Union @@ -8,6 +9,14 @@ import requests from tqdm import tqdm +def import_package(name, package=None): + try: + module = importlib.import_module(name, package=package) + return module + except ModuleNotFoundError: + return None + + def download_file(url: str, save_path: Union[str, Path], logger=None): if not Path(save_path).parent.exists(): Path(save_path).parent.mkdir(parents=True, exist_ok=True) diff --git a/python/setup.py b/python/setup.py index 6e691bc..fd70ed8 100644 --- a/python/setup.py +++ b/python/setup.py @@ -58,7 +58,7 @@ setuptools.setup( project_urls=project_urls, license="Apache-2.0", include_package_data=True, - install_requires=read_txt("requirements_ort.txt"), + install_requires=read_txt("requirements.txt"), package_dir={"": MODULE_NAME}, packages=setuptools.find_namespace_packages(where=MODULE_NAME), package_data={"": ["*.onnx", "*.yaml"]},