diff --git a/python/rapidocr/config.yaml b/python/rapidocr/config.yaml index 14eaa65..7f3c019 100644 --- a/python/rapidocr/config.yaml +++ b/python/rapidocr/config.yaml @@ -81,6 +81,8 @@ EngineConfig: npu_ep_cfg: device_id: 0 + mnn: {} + Det: engine_type: "onnxruntime" lang_type: "ch" diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index 82c14f0..aa39e4b 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/python/rapidocr/inference_engine/base.py @@ -1,4 +1,4 @@ -# -*- encoding: utf-8 -*- + # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com import abc @@ -53,6 +53,14 @@ def get_engine(engine_type: EngineType): return TorchInferSession + if engine_type == EngineType.MNN: + if not import_package("MNN"): + raise ImportError("MNN is not installed") + + from .mnn import MNNInferSession + + return MNNInferSession + raise ValueError(f"Unsupported engine: {engine_type.value}") diff --git a/python/rapidocr/inference_engine/mnn.py b/python/rapidocr/inference_engine/mnn.py new file mode 100644 index 0000000..45034ea --- /dev/null +++ b/python/rapidocr/inference_engine/mnn.py @@ -0,0 +1,79 @@ +# -*- encoding: utf-8 -*- +import traceback +from pathlib import Path + +import MNN +import numpy as np +from omegaconf import DictConfig + +from ..utils.download_file import DownloadFile, DownloadFileInput +from ..utils.log import logger +from .base import FileInfo, InferSession + + +class MNNInferSession(InferSession): + def __init__(self, cfg: DictConfig): + super().__init__(cfg) + + model_path = cfg.get("model_path", None) + if model_path is None: + model_info = self.get_model_url( + FileInfo( + engine_type=cfg.engine_type, + ocr_version=cfg.ocr_version, + task_type=cfg.task_type, + lang_type=cfg.lang_type, + model_type=cfg.model_type, + ) + ) + model_path = self.DEFAULT_MODEL_PATH / Path(model_info["model_dir"]).name + DownloadFile.run( + DownloadFileInput( + file_url=model_info["model_dir"], + sha256=model_info["SHA256"], + save_path=model_path, + logger=logger, + ) + ) + + model_path = Path(model_path) + self._verify_model(model_path) + self.interpreter = MNN.Interpreter(str(model_path)) + self.session = self.interpreter.createSession() + self.input_tensor = self.interpreter.getSessionInput(self.session) + + def __call__(self, img: np.ndarray) -> np.ndarray: + try: + self.interpreter.resizeTensor(self.input_tensor, img.shape) + self.interpreter.resizeSession(self.session) + + tmp = MNN.Tensor( + img.shape, + MNN.Halide_Type_Float, + img, + MNN.Tensor_DimensionType_Caffe, + ) + self.input_tensor.copyFromHostTensor(tmp) + self.interpreter.runSession(self.session) + + output = self.interpreter.getSessionOutput(self.session) + out_shape = output.getShape() + out_tensor = MNN.Tensor( + out_shape, + MNN.Halide_Type_Float, + MNN.Tensor_DimensionType_Caffe + ) + output.copyToHostTensor(out_tensor) + + return np.array(out_tensor.getData()).reshape(out_shape) + + except Exception as e: + error_info = traceback.format_exc() + raise MNNError(f"MNN inference failed: {error_info}") from e + + def have_key(self, key: str = "character") -> bool: + return False + + +class MNNError(Exception): + pass diff --git a/python/rapidocr/utils/typings.py b/python/rapidocr/utils/typings.py index 67aa2db..b0ec1ad 100644 --- a/python/rapidocr/utils/typings.py +++ b/python/rapidocr/utils/typings.py @@ -53,6 +53,7 @@ class EngineType(Enum): OPENVINO = "openvino" PADDLE = "paddle" TORCH = "torch" + MNN = "mnn" class ModelType(Enum):