Optim rapid_structure ort infer

python/demo.py

@@ -92,7 +92,6 @@ if __name__ == '__main__':
 
     image_path = 'tests/test_files/ch_en_num.jpg'
     img = cv2.imread(image_path)
-    img = np.zeros([640, 640, 3])
     result, elapse_list = rapid_ocr(img)
     print(result)
     print(elapse_list)

python/rapid_structure/demo.py

@@ -80,6 +80,6 @@ def demo_orientation():
 
 
 if __name__ == '__main__':
-    # demo_layout()
+    demo_layout()
     # demo_table()
-    demo_orientation()
+    # demo_orientation()

python/rapid_structure/rapid_layout/rapid_layout.py

@@ -33,10 +33,8 @@ class RapidLayout():
             model_path = str(root_dir / 'models' / 'layout_cdla.onnx')
         config['model_path'] = model_path
 
-        session_instance = OrtInferSession(config)
-        self.predictor = session_instance.session
-        self.input_name = session_instance.get_input_name()
-        labels = session_instance.get_metadata()['character'].splitlines()
+        self.session = OrtInferSession(config)
+        labels = self.session.get_metadata()['character'].splitlines()
 
         self.preprocess_op = create_operators(config['pre_process'])
         self.postprocess_op = PicoDetPostProcess(labels,
@@ -57,16 +55,14 @@ class RapidLayout():
         preds, elapse = 0, 1
         starttime = time.time()
 
-        input_dict = {self.input_name: img}
-        preds_onnx = self.predictor.run(None, input_dict)
+        preds = self.session(img)[0]
 
-        np_score_list, np_boxes_list = [], []
-        num_outs = int(len(preds_onnx) / 2)
+        score_list, boxes_list = [], []
+        num_outs = int(len(preds) / 2)
         for out_idx in range(num_outs):
-            np_score_list.append(preds_onnx[out_idx])
-            np_boxes_list.append(preds_onnx[out_idx + num_outs])
-        preds = dict(boxes=np_score_list, boxes_num=np_boxes_list)
-
+            score_list.append(preds[out_idx])
+            boxes_list.append(preds[out_idx + num_outs])
+        preds = dict(boxes=score_list, boxes_num=boxes_list)
         post_preds = self.postprocess_op(ori_im, img, preds)
         elapse = time.time() - starttime
         return post_preds, elapse

python/rapid_structure/rapid_layout/utils.py

@@ -3,6 +3,7 @@
 # @Contact: liekkaskono@163.com
 import copy
 import warnings
+from pathlib import Path
 
 import yaml
 import cv2
@@ -29,6 +30,7 @@ class OrtInferSession():
             EP_list = [(cuda_ep, config[cuda_ep])]
         EP_list.append((cpu_ep, cpu_provider_options))
 
+        self._verify_model(config['model_path'])
         self.session = InferenceSession(config['model_path'],
                                         sess_options=sess_opt,
                                         providers=EP_list)
@@ -40,16 +42,35 @@ 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_metadata(self):
         meta_dict = self.session.get_modelmeta().custom_metadata_map
         return meta_dict
 
+    @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 ONNXRuntimeError(Exception):
+    pass
+
 
 def transform(data, ops=None):
     """ transform """

python/rapid_structure/rapid_orientation/rapid_orientation.py

@@ -35,10 +35,8 @@ class RapidOrientation():
             model_path = str(root_dir / 'models' / 'rapid_orientation.onnx')
         config['model_path'] = model_path
 
-        session_instance = OrtInferSession(config)
-        self.predictor = session_instance.session
-        self.input_name = session_instance.get_input_name()
-        self.labels = session_instance.get_metadata()['character'].splitlines()
+        self.session = OrtInferSession(config)
+        self.labels = self.session.get_metadata()['character'].splitlines()
 
         self.preprocess_ops = create_operators(config["PreProcess"])
 
@@ -46,11 +44,10 @@ class RapidOrientation():
         s = time.time()
         for ops in self.preprocess_ops:
             images = ops(images)
-        image = np.array(images)
-        image = image[np.newaxis, ...]
+        image = np.array(images)[None, ...]
+
+        pred_output = self.session(image)[0]
 
-        input_dict = {self.input_name: image}
-        pred_output = self.predictor.run(None, input_feed=input_dict)[0]
         pred_output = pred_output.squeeze()
         pred_idx = np.argmax(pred_output)
         pred_txt = self.labels[pred_idx]

python/rapid_structure/rapid_orientation/utils.py

@@ -2,6 +2,7 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 import importlib
+from pathlib import Path
 import warnings
 
 import cv2
@@ -27,6 +28,7 @@ class OrtInferSession():
             EP_list = [(cuda_ep, config[cuda_ep])]
         EP_list.append((cpu_ep, cpu_provider_options))
 
+        self._verify_model(config['model_path'])
         self.session = InferenceSession(config['model_path'],
                                         sess_options=sess_opt,
                                         providers=EP_list)
@@ -42,16 +44,35 @@ class OrtInferSession():
                           f'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_metadata(self):
         meta_dict = self.session.get_modelmeta().custom_metadata_map
         return meta_dict
 
+    @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 ONNXRuntimeError(Exception):
+    pass
+
 
 def create_operators(params):
     """

python/rapid_structure/rapid_table/table_structure/table_structure.py

@@ -21,11 +21,9 @@ class TableStructurer():
     def __init__(self, model_path: str):
         self.preprocess_op = TablePreprocess()
 
-        session_instance = OrtInferSession(model_path)
-        self.predictor = session_instance.session
-        self.input_name = session_instance.get_input_name()
+        self.session = OrtInferSession(model_path)
 
-        self.character = session_instance.get_metadata()
+        self.character = self.session.get_metadata()
         self.postprocess_op = TableLabelDecode(self.character)
 
     def __call__(self, img):
@@ -38,8 +36,7 @@ class TableStructurer():
         img = np.expand_dims(img, axis=0)
         img = img.copy()
 
-        input_dict = {self.input_name: img}
-        outputs = self.predictor.run(None, input_dict)
+        outputs = self.session(img)
 
         preds = {
             'loc_preds': outputs[0],

python/rapid_structure/rapid_table/table_structure/utils.py

@@ -24,7 +24,7 @@ from onnxruntime import (GraphOptimizationLevel, InferenceSession,
 
 class OrtInferSession():
     def __init__(self, onnx_path: str):
-        self.__verify_model(onnx_path)
+        self._verify_model(onnx_path)
 
         sess_opt = SessionOptions()
         sess_opt.log_severity_level = 4
@@ -35,15 +35,23 @@ class OrtInferSession():
         cpu_provider_options = {"arena_extend_strategy": "kSameAsRequested",}
         EP_list = [(cpu_ep, cpu_provider_options)]
 
+        self._verify_model(onnx_path)
         self.session = InferenceSession(onnx_path,
                                         sess_options=sess_opt,
                                         providers=EP_list)
 
-    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_metadata(self, key: str = 'character') -> list:
         meta_dict = self.session.get_modelmeta().custom_metadata_map
@@ -51,7 +59,7 @@ class OrtInferSession():
         return content_list
 
     @staticmethod
-    def __verify_model(model_path):
+    def _verify_model(model_path):
         model_path = Path(model_path)
         if not model_path.exists():
             raise FileNotFoundError(f'{model_path} does not exists.')
@@ -59,6 +67,10 @@ class OrtInferSession():
             raise FileExistsError(f'{model_path} is not a file.')
 
 
+class ONNXRuntimeError(Exception):
+    pass
+
+
 class TableLabelDecode():
     """  """
     def __init__(self,

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py

@@ -55,7 +55,7 @@ class OrtInferSession():
                           RuntimeWarning)
 
     def __call__(self, input_content: np.ndarray) -> np.ndarray:
-        input_dict = dict(zip(self.get_input_names(), input_content))
+        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:

python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py

@@ -35,7 +35,7 @@ class TextDetector():
 
         self.infer = OrtInferSession(config)
 
-    def __call__(self, img: np.ndarry):
+    def __call__(self, img):
         if img is None:
             raise ValueError('img is None')
 

python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py

@@ -65,13 +65,13 @@ class OrtInferSession():
                           RuntimeWarning)
 
     def __call__(self, input_content: np.ndarray) -> np.ndarray:
-        input_dict = dict(zip(self.get_input_names(), input_content))
+        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_input_names(self, )
+    def get_input_names(self,):
         return [v.name for v in self.session.get_inputs()]
 
     def get_output_names(self,):

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py

@@ -45,7 +45,7 @@ class OrtInferSession():
                           RuntimeWarning)
 
     def __call__(self, input_content: np.ndarray) -> np.ndarray:
-        input_dict = dict(zip(self.get_input_names(), input_content))
+        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:

python/tests/test_cls.py

@@ -8,14 +8,14 @@ from base_module import BaseModule
 
 @pytest.mark.parametrize(
     'package_name',
-    ['rapidocr_onnxruntime', 'rapidocr_openvino']
+    [('rapidocr_onnxruntime'),
+     ('rapidocr_openvino')]
 )
 def test_cls(package_name: str):
     module_name = 'ch_ppocr_v2_cls'
     class_name = 'TextClassifier'
 
     base = BaseModule(package_name=package_name)
-
     TextClassifier = base.init_module(module_name, class_name)
 
     yaml_path = base.package_dir / module_name / 'config.yaml'