diff --git a/python/rapidocr/ch_ppocr_cls/main.py b/python/rapidocr/ch_ppocr_cls/main.py index 7b5b207..851239d 100644 --- a/python/rapidocr/ch_ppocr_cls/main.py +++ b/python/rapidocr/ch_ppocr_cls/main.py @@ -19,7 +19,7 @@ from typing import Any, Dict, List, Union import cv2 import numpy as np -from rapidocr.utils import OrtInferSession +from rapidocr.inference_engine import OrtInferSession from .utils import ClsPostProcess, TextClsOutput diff --git a/python/rapidocr/ch_ppocr_det/main.py b/python/rapidocr/ch_ppocr_det/main.py index a6cce33..63e633d 100644 --- a/python/rapidocr/ch_ppocr_det/main.py +++ b/python/rapidocr/ch_ppocr_det/main.py @@ -19,7 +19,7 @@ from typing import Any, Dict, List import numpy as np -from rapidocr.utils import OrtInferSession +from rapidocr.inference_engine import OrtInferSession from .utils import DBPostProcess, DetPreProcess, TextDetOutput diff --git a/python/rapidocr/ch_ppocr_rec/main.py b/python/rapidocr/ch_ppocr_rec/main.py index 44ea52b..26e7a81 100644 --- a/python/rapidocr/ch_ppocr_rec/main.py +++ b/python/rapidocr/ch_ppocr_rec/main.py @@ -18,7 +18,7 @@ from typing import Any, Dict import cv2 import numpy as np -from rapidocr.utils import OrtInferSession +from rapidocr.inference_engine import OrtInferSession from .utils import CTCLabelDecode, TextRecInput, TextRecOutput diff --git a/python/rapidocr/inference_engine/__init__.py b/python/rapidocr/inference_engine/__init__.py index 0ecdd4f..9b75723 100644 --- a/python/rapidocr/inference_engine/__init__.py +++ b/python/rapidocr/inference_engine/__init__.py @@ -1,3 +1,4 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from .onnxruntime import OrtInferSession diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py new file mode 100644 index 0000000..1f0f526 --- /dev/null +++ b/python/rapidocr/inference_engine/base.py @@ -0,0 +1,30 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com +import abc +from pathlib import Path +from typing import Union + +import numpy as np + + +class InferSession(abc.ABC): + @abc.abstractmethod + def __init__(self, config): + pass + + @abc.abstractmethod + def __call__(self, input_content: np.ndarray) -> np.ndarray: + pass + + @staticmethod + def _verify_model(model_path: Union[str, Path, None]): + if model_path is None: + raise ValueError("model_path is None!") + + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f"{model_path} does not exists.") + + if not model_path.is_file(): + raise FileExistsError(f"{model_path} is not a file.") diff --git a/python/rapidocr/inference_engine/onnxruntime.py b/python/rapidocr/inference_engine/onnxruntime.py index 09f688d..45967a0 100644 --- a/python/rapidocr/inference_engine/onnxruntime.py +++ b/python/rapidocr/inference_engine/onnxruntime.py @@ -5,8 +5,7 @@ import os import platform import traceback from enum import Enum -from pathlib import Path -from typing import Any, Dict, List, Tuple, Union +from typing import Any, Dict, List, Tuple import numpy as np from onnxruntime import ( @@ -18,6 +17,7 @@ from onnxruntime import ( ) from ..utils.logger import get_logger +from .base import InferSession class EP(Enum): @@ -26,7 +26,7 @@ class EP(Enum): DIRECTML_EP = "DmlExecutionProvider" -class OrtInferSession: +class OrtInferSession(InferSession): def __init__(self, config: Dict[str, Any]): self.logger = get_logger("OrtInferSession") @@ -214,18 +214,6 @@ class OrtInferSession: return True return False - @staticmethod - def _verify_model(model_path: Union[str, Path, None]): - if model_path is None: - raise ValueError("model_path is None!") - - model_path = Path(model_path) - if not model_path.exists(): - raise FileNotFoundError(f"{model_path} does not exists.") - - if not model_path.is_file(): - raise FileExistsError(f"{model_path} is not a file.") - class ONNXRuntimeError(Exception): pass diff --git a/python/rapidocr/inference_engine/torch.py b/python/rapidocr/inference_engine/torch.py index 0ecdd4f..d59202a 100644 --- a/python/rapidocr/inference_engine/torch.py +++ b/python/rapidocr/inference_engine/torch.py @@ -1,3 +1,65 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from pathlib import Path +from typing import Dict, Optional, Union + +import numpy as np +import torch +import yaml + +root_dir = Path(__file__).resolve().parent.parent +DEFAULT_CFG_PATH = root_dir / "arch_config.yaml" + + +def read_yaml(yaml_path: Union[str, Path]) -> Dict[str, Dict]: + with open(yaml_path, "rb") as f: + data = yaml.load(f, Loader=yaml.Loader) + return data + + +from rapidocr_torch.modeling.architectures.base_model import BaseModel + +from ..utils.logger import get_logger + + +class TorchInferSession: + def __init__(self, config, mode: Optional[str] = None) -> None: + all_arch_config = read_yaml(DEFAULT_CFG_PATH) + + self.logger = get_logger("TorchInferSession") + self.mode = mode + model_path = Path(config["model_path"]) + self._verify_model(model_path) + file_name = model_path.stem + if file_name not in all_arch_config: + raise ValueError(f"architecture {file_name} is not in config.yaml") + arch_config = all_arch_config[file_name] + self.predictor = BaseModel(arch_config) + self.predictor.load_state_dict(torch.load(model_path, weights_only=True)) + self.predictor.eval() + self.use_gpu = False + if config["use_cuda"]: + self.predictor.cuda() + self.use_gpu = True + + def __call__(self, img: np.ndarray): + with torch.no_grad(): + inp = torch.from_numpy(img) + if self.use_gpu: + inp = inp.cuda() + # 适配跟onnx对齐取值逻辑 + outputs = self.predictor(inp).unsqueeze(0) + return outputs.cpu().numpy() + + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f"{model_path} does not exists.") + if not model_path.is_file(): + raise FileExistsError(f"{model_path} is not a file.") + + +class TorchInferError(Exception): + pass diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index 43f605b..c39c694 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -44,8 +44,8 @@ class RapidOCR: updater = UpdateParameters() config = updater(config, **kwargs) - det_lang, rec_lang = parse_lang(config.Global.lang) # 根据选定的语言加载对应的模型 + det_lang, rec_lang = parse_lang(config.Global.lang) self.print_verbose = config.Global.print_verbose self.text_score = config.Global.text_score diff --git a/python/rapidocr/utils/__init__.py b/python/rapidocr/utils/__init__.py index dea0484..7535b42 100644 --- a/python/rapidocr/utils/__init__.py +++ b/python/rapidocr/utils/__init__.py @@ -1,7 +1,6 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -from ..inference_engine.onnxruntime import OrtInferSession from .load_image import LoadImage, LoadImageError from .logger import get_logger from .parse_parameters import UpdateParameters, init_args, parse_lang, update_model_path