Optim ort infer

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py

@@ -33,9 +33,7 @@ class TextClassifier():
         self.cls_thresh = config['cls_thresh']
         self.postprocess_op = ClsPostProcess(config['label_list'])
 
-        session_instance = OrtInferSession(config)
-        self.session = session_instance.session
-        self.input_name = session_instance.get_input_name()
+        self.infer = OrtInferSession(config)
 
     def __call__(self, img_list: List[np.ndarray]):
         if isinstance(img_list, np.ndarray):
@@ -64,8 +62,7 @@ class TextClassifier():
             norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
 
             starttime = time.time()
-            onnx_inputs = {self.input_name: norm_img_batch}
-            prob_out = self.session.run(None, onnx_inputs)[0]
+            prob_out = self.infer(norm_img_batch)[0]
             cls_result = self.postprocess_op(prob_out)
             elapse += time.time() - starttime
 

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py

@@ -15,6 +15,7 @@ from pathlib import Path
 import warnings
 
 import yaml
+import numpy as np
 from onnxruntime import (get_available_providers, get_device,
                          SessionOptions, InferenceSession,
                          GraphOptimizationLevel)
@@ -53,11 +54,18 @@ class OrtInferSession():
                           'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
                           RuntimeWarning)
 
-    def get_input_name(self, input_idx=0):
-        return self.session.get_inputs()[input_idx].name
+    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)
+        except Exception as e:
+            raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
 
-    def get_output_name(self, output_idx=0):
-        return self.session.get_outputs()[output_idx].name
+    def get_input_names(self, ):
+        return [v.name for v in self.session.get_inputs()]
+
+    def get_output_names(self,):
+        return [v.name for v in self.session.get_outputs()]
 
     @staticmethod
     def _verify_model(model_path):
@@ -68,6 +76,10 @@ class OrtInferSession():
             raise FileExistsError(f'{model_path} is not a file.')
 
 
+class ONNXRuntimeError(Exception):
+    pass
+
+
 def read_yaml(yaml_path):
     with open(yaml_path, 'rb') as f:
         data = yaml.load(f, Loader=yaml.Loader)

python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py

@@ -33,11 +33,9 @@ class TextDetector():
         self.preprocess_op = create_operators(config['pre_process'])
         self.postprocess_op = DBPostProcess(**config['post_process'])
 
-        session_instance = OrtInferSession(config)
-        self.session = session_instance.session
-        self.input_name = session_instance.get_input_name()
+        self.infer = OrtInferSession(config)
 
-    def __call__(self, img):
+    def __call__(self, img: np.ndarry):
         if img is None:
             raise ValueError('img is None')
 
@@ -53,9 +51,8 @@ class TextDetector():
         shape_list = np.expand_dims(shape_list, axis=0)
 
         starttime = time.time()
-        preds = self.session.run(None, {self.input_name: img})
-
-        post_result = self.postprocess_op(preds[0], shape_list)
+        preds = self.infer(img)[0]
+        post_result = self.postprocess_op(preds, shape_list)
 
         dt_boxes = post_result[0]['points']
         dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im_shape)

python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py

@@ -64,11 +64,18 @@ class OrtInferSession():
                           'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
                           RuntimeWarning)
 
-    def get_input_name(self, input_idx=0):
-        return self.session.get_inputs()[input_idx].name
+    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)
+        except Exception as e:
+            raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
 
-    def get_output_name(self, output_idx=0):
-        return self.session.get_outputs()[output_idx].name
+    def get_input_names(self, )
+        return [v.name for v in self.session.get_inputs()]
+
+    def get_output_names(self,):
+        return [v.name for v in self.session.get_outputs()]
 
     @staticmethod
     def _verify_model(model_path):
@@ -79,6 +86,10 @@ class OrtInferSession():
             raise FileExistsError(f'{model_path} is not a file.')
 
 
+class ONNXRuntimeError(Exception):
+    pass
+
+
 def read_yaml(yaml_path):
     with open(yaml_path, 'rb') as f:
         data = yaml.load(f, Loader=yaml.Loader)

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py

@@ -27,12 +27,10 @@ except:
 
 class TextRecognizer():
     def __init__(self, config):
-        session_instance = OrtInferSession(config)
-        self.session = session_instance.session
-        self.input_name = session_instance.get_input_name()
+        self.session = OrtInferSession(config)
 
-        if session_instance.have_key():
-            self.character_dict_path = session_instance.get_character_list()
+        if self.session.have_key():
+            self.character_dict_path = self.session.get_character_list()
         else:
             self.character_dict_path = config.get('keys_path', None)
         self.postprocess_op = CTCLabelDecode(self.character_dict_path)
@@ -71,8 +69,7 @@ class TextRecognizer():
             norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
 
             starttime = time.time()
-            onnx_inputs = {self.input_name: norm_img_batch}
-            preds = self.session.run(None, onnx_inputs)[0]
+            preds = self.session(norm_img_batch)[0]
             rec_result = self.postprocess_op(preds)
 
             for rno in range(len(rec_result)):

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py

@@ -44,11 +44,18 @@ class OrtInferSession():
                           'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
                           RuntimeWarning)
 
-    def get_input_name(self, input_idx=0):
-        return self.session.get_inputs()[input_idx].name
+    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)
+        except Exception as e:
+            raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e
 
-    def get_output_name(self, output_idx=0):
-        return self.session.get_outputs()[output_idx].name
+    def get_input_names(self, ):
+        return [v.name for v in self.session.get_inputs()]
+
+    def get_output_names(self,):
+        return [v.name for v in self.session.get_outputs()]
 
     def get_character_list(self, key: str = 'character'):
         return self.meta_dict[key].splitlines()
@@ -68,6 +75,10 @@ class OrtInferSession():
             raise FileExistsError(f'{model_path} is not a file.')
 
 
+class ONNXRuntimeError(Exception):
+    pass
+
+
 def read_yaml(yaml_path):
     with open(yaml_path, 'rb') as f:
         data = yaml.load(f, Loader=yaml.Loader)

python/tests/test_all_ort.py

@@ -15,7 +15,7 @@ sys.path.append(str(root_dir))
 from rapidocr_onnxruntime import RapidOCR
 
 rapid_ocr = RapidOCR()
-tests_dir = root_dir / 'test_files'
+tests_dir = root_dir / 'tests' / 'test_files'
 
 
 def test_normal():

python/tests/test_all_vino.py

@@ -16,7 +16,7 @@ from rapidocr_openvino import RapidOCR
 
 rapid_ocr = RapidOCR()
 
-tests_dir = root_dir / 'test_files'
+tests_dir = root_dir / 'tests' / 'test_files'
 
 
 def test_normal():