chore: update files

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.inference_engine import OrtInferSession
+from rapidocr.inference_engine import get_engine
 
 from .utils import ClsPostProcess, TextClsOutput
 
@@ -31,7 +31,7 @@ class TextClassifier:
         self.cls_thresh = config["cls_thresh"]
         self.postprocess_op = ClsPostProcess(config["label_list"])
 
-        self.infer = OrtInferSession(config)
+        self.session = get_engine(config.engine_name)(config)
 
     def __call__(self, img_list: Union[np.ndarray, List[np.ndarray]]) -> TextClsOutput:
         if isinstance(img_list, np.ndarray):
@@ -60,7 +60,7 @@ class TextClassifier:
             norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
 
             starttime = time.time()
-            prob_out = self.infer(norm_img_batch)[0]
+            prob_out = self.session(norm_img_batch)
             cls_result = self.postprocess_op(prob_out)
             elapse += time.time() - starttime
 

python/rapidocr/ch_ppocr_det/main.py

@@ -19,7 +19,7 @@ from typing import Any, Dict, List
 
 import numpy as np
 
-from rapidocr.inference_engine import OrtInferSession
+from rapidocr.inference_engine import get_engine
 
 from .utils import DBPostProcess, DetPreProcess, TextDetOutput
 
@@ -42,7 +42,7 @@ class TextDetector:
         }
         self.postprocess_op = DBPostProcess(**post_process)
 
-        self.infer = OrtInferSession(config)
+        self.session = get_engine(config.engine_name)(config)
 
     def __call__(self, img: np.ndarray) -> TextDetOutput:
         start_time = time.perf_counter()
@@ -56,7 +56,7 @@ class TextDetector:
         if prepro_img is None:
             return TextDetOutput()
 
-        preds = self.infer(prepro_img)[0]
+        preds = self.session(prepro_img)
         boxes, scores = self.postprocess_op(preds, ori_img_shape)
         if len(boxes) < 1:
             return TextDetOutput()

python/rapidocr/ch_ppocr_rec/main.py

@@ -13,27 +13,31 @@
 # limitations under the License.
 import math
 import time
+from pathlib import Path
 from typing import Any, Dict
 
 import cv2
 import numpy as np
 
-from rapidocr.inference_engine import OrtInferSession
+from rapidocr.inference_engine import get_engine
 
 from .utils import CTCLabelDecode, TextRecInput, TextRecOutput
 
+DEFAULT_DICT_PATH = str(Path(__file__).parent.parent / "models" / "ppocr_keys_v1.txt")
+
 
 class TextRecognizer:
     def __init__(self, config: Dict[str, Any]):
-        self.session = OrtInferSession(config)
+        self.session = get_engine(config.engine_name)(config)
 
         character = None
         if self.session.have_key():
             character = self.session.get_character_list()
 
-        character_path = config.get("rec_keys_path", None)
+        dict_path = config.get("rec_keys_path", None)
+        character_dict_path = dict_path if dict_path else DEFAULT_DICT_PATH
         self.postprocess_op = CTCLabelDecode(
-            character=character, character_path=character_path
+            character=character, character_path=character_dict_path
         )
 
         self.rec_batch_num = config["rec_batch_num"]
@@ -73,7 +77,7 @@ class TextRecognizer:
             norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
 
             start_time = time.perf_counter()
-            preds = self.session(norm_img_batch)[0]
+            preds = self.session(norm_img_batch)
             line_results, word_results = self.postprocess_op(
                 preds,
                 return_word_box,

python/rapidocr/config.yaml

@@ -14,6 +14,11 @@ Global:
     intra_op_num_threads: &intra_nums -1
     inter_op_num_threads: &inter_nums -1
 
+    with_onnx: false
+    with_openvino: true
+    with_paddle: false
+    with_torch: false
+
 Det:
     intra_op_num_threads: *intra_nums
     inter_op_num_threads: *inter_nums

python/rapidocr/inference_engine/__init__.py

@@ -1,4 +1,55 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+from enum import Enum
+
+from omegaconf import DictConfig
+
 from .onnxruntime import OrtInferSession
+from .openvino import OpenVINOInferSession
+from .paddlepaddle import PaddleInferSession
+from .torch import TorchInferSession
+
+
+class Engine(Enum):
+    ONNXRUNTIME = "onnxruntime"
+    OPENVINO = "openvino"
+    PADDLE = "paddlepaddle"
+    TORCH = "torch"
+
+
+def get_engine(engine_name: str):
+    if engine_name == Engine.ONNXRUNTIME.value:
+        return OrtInferSession
+
+    if engine_name == Engine.OPENVINO.value:
+        return OpenVINOInferSession
+
+    if engine_name == Engine.PADDLE.value:
+        return PaddleInferSession
+
+    if engine_name == Engine.TORCH.value:
+        return TorchInferSession
+
+    raise ValueError(f"Unsupported engine: {engine_name}")
+
+
+def get_engine_name(config: DictConfig) -> str:
+    with_onnx = config.Global.with_onnx
+    with_openvino = config.Global.with_openvino
+    with_paddle = config.Global.with_paddle
+    with_torch = config.Global.with_torch
+
+    if with_onnx:
+        return Engine.ONNXRUNTIME.value
+
+    if with_openvino:
+        return Engine.OPENVINO.value
+
+    if with_paddle:
+        return Engine.PADDLE.value
+
+    if with_torch:
+        return Engine.TORCH.value
+
+    return Engine.ONNXRUNTIME.value

python/rapidocr/inference_engine/base.py

@@ -28,3 +28,7 @@ class InferSession(abc.ABC):
 
         if not model_path.is_file():
             raise FileExistsError(f"{model_path} is not a file.")
+
+    @abc.abstractmethod
+    def have_key(self, key: str = "character") -> bool:
+        pass

python/rapidocr/inference_engine/onnxruntime.py

@@ -193,7 +193,7 @@ class OrtInferSession(InferSession):
     def __call__(self, input_content: np.ndarray) -> np.ndarray:
         input_dict = dict(zip(self.get_input_names(), [input_content]))
         try:
-            return self.session.run(self.get_output_names(), input_dict)
+            return self.session.run(self.get_output_names(), input_dict)[0]
         except Exception as e:
             error_info = traceback.format_exc()
             raise ONNXRuntimeError(error_info) from e

python/rapidocr/inference_engine/openvino.py

@@ -3,13 +3,14 @@
 # @Contact: liekkaskono@163.com
 import os
 import traceback
-from pathlib import Path
 
 import numpy as np
 from openvino.runtime import Core
 
+from .base import InferSession
 
-class OpenVINOInferSession:
+
+class OpenVINOInferSession(InferSession):
     def __init__(self, config):
         core = Core()
 
@@ -32,13 +33,8 @@ class OpenVINOInferSession:
             error_info = traceback.format_exc()
             raise OpenVIONError(error_info) from e
 
-    @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.")
+    def have_key(self, key: str = "character") -> bool:
+        return False
 
 
 class OpenVIONError(Exception):

python/rapidocr/main.py

@@ -13,6 +13,7 @@ from .cal_rec_boxes import CalRecBoxes
 from .ch_ppocr_cls import TextClassifier, TextClsOutput
 from .ch_ppocr_det import TextDetector, TextDetOutput
 from .ch_ppocr_rec import TextRecInput, TextRecognizer, TextRecOutput
+from .inference_engine import get_engine_name
 from .utils import (
     LoadImage,
     RapidOCROutput,
@@ -44,6 +45,8 @@ class RapidOCR:
             updater = UpdateParameters()
             config = updater(config, **kwargs)
 
+        engine_name = get_engine_name(config)
+
         # 根据选定的语言加载对应的模型
         det_lang, rec_lang = parse_lang(config.Global.lang)
 
@@ -54,13 +57,16 @@ class RapidOCR:
 
         self.use_det = config.Global.use_det
         config.Det.lang = det_lang
+        config.Det.engine_name = engine_name
         self.text_det = TextDetector(config.Det)
 
         self.use_cls = config.Global.use_cls
+        config.Cls.engine_name = engine_name
         self.text_cls = TextClassifier(config.Cls)
 
         self.use_rec = config.Global.use_rec
         config.Rec.lang = rec_lang
+        config.Rec.engine_name = engine_name
         self.text_rec = TextRecognizer(config.Rec)
 
         self.load_img = LoadImage()
@@ -303,6 +309,10 @@ class RapidOCR:
         ocr_res.scores = tuple(filter_scores)
         return ocr_res
 
+    def export_config(self, save_path: Union[Path, str]) -> None:
+        with open(save_path, "w", encoding="utf-8") as f:
+            OmegaConf.save(self.config, f)
+
 
 def main():
     args = init_args()