Optim rapidocr_openvino

python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py

@@ -22,10 +22,10 @@ import numpy as np
 
 try:
     from .utils import (DBPostProcess, create_operators,
-                        draw_text_det_res, transform, read_yaml, OrtInferSession)
+                        transform, read_yaml, OrtInferSession)
 except:
     from utils import (DBPostProcess, create_operators,
-                       draw_text_det_res, transform, read_yaml, OrtInferSession)
+                       transform, read_yaml, OrtInferSession)
 
 
 class TextDetector(object):
@@ -120,6 +120,8 @@ if __name__ == "__main__":
 
     img = cv2.imread(args.image_path)
     dt_boxes, elapse = text_detector(img)
+
+    from utils import draw_text_det_res
     src_im = draw_text_det_res(dt_boxes, args.image_path)
     cv2.imwrite('det_results.jpg', src_im)
     print('The det_results.jpg has been saved in the current directory.')

python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py

@@ -113,4 +113,4 @@ if __name__ == "__main__":
 
     img = cv2.imread(args.image_path)
     rec_res, predict_time = text_recognizer(img)
-    print(f'识别结果: {rec_res}\t cost: {predict_time}s')
+    print(f'rec result: {rec_res}\t cost: {predict_time}s')

python/rapidocr_openvino/ch_ppocr_v2_cls/config.yaml

@@ -1,5 +1,13 @@
 model_path: resources/models/ch_ppocr_mobile_v2.0_cls_infer.onnx
 
+use_cuda: false
+# Details of the params: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html
+CUDAExecutionProvider:
+    device_id: 0
+    arena_extend_strategy: kNextPowerOfTwo
+    cudnn_conv_algo_search: EXHAUSTIVE
+    do_copy_in_default_stream: true
+
 cls_image_shape: [3, 48, 192]
 cls_batch_num: 6
 cls_thresh: 0.9

python/rapidocr_openvino/ch_ppocr_v2_cls/text_cls.py

@@ -19,12 +19,11 @@ from typing import List
 
 import cv2
 import numpy as np
-import onnxruntime as ort
 
 try:
-    from .utils import ClsPostProcess, read_yaml
+    from .utils import ClsPostProcess, read_yaml, OrtInferSession
 except:
-    from utils import ClsPostProcess, read_yaml
+    from utils import ClsPostProcess, read_yaml, OrtInferSession
 
 
 class TextClassifier(object):
@@ -34,11 +33,9 @@ class TextClassifier(object):
         self.cls_thresh = config['cls_thresh']
         self.postprocess_op = ClsPostProcess(config['label_list'])
 
-        sess_opt = ort.SessionOptions()
-        sess_opt.log_severity_level = 4
-        sess_opt.enable_cpu_mem_arena = False
-        self.session = ort.InferenceSession(config['model_path'], sess_opt)
-        self.input_name = self.session.get_inputs()[0].name
+        session_instance = OrtInferSession(config)
+        self.session = session_instance.session
+        self.input_name = session_instance.get_input_name()
 
     def __call__(self, img_list: List[np.ndarray]):
         if isinstance(img_list, np.ndarray):
@@ -122,4 +119,4 @@ if __name__ == "__main__":
     img = cv2.imread(args.image_path)
     img_list, cls_res, predict_time = text_classifier(img)
     for ino in range(len(img_list)):
-        print(f"分类结果:{cls_res[ino]}")
+        print(f"cls result:{cls_res[ino]}")

python/rapidocr_openvino/ch_ppocr_v2_cls/utils.py

@@ -11,7 +11,44 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import warnings
+
 import yaml
+from onnxruntime import (get_available_providers, get_device,
+                         SessionOptions, InferenceSession)
+
+
+class OrtInferSession(object):
+    def __init__(self, config):
+        sess_opt = SessionOptions()
+        sess_opt.log_severity_level = 4
+        sess_opt.enable_cpu_mem_arena = False
+
+        cuda_ep = 'CUDAExecutionProvider'
+        cpu_ep = 'CPUExecutionProvider'
+
+        EP_list = []
+        if config['use_cuda'] and get_device() == 'GPU' \
+                and cuda_ep in get_available_providers():
+            EP_list = [(cuda_ep, config[cuda_ep])]
+        EP_list.append(cpu_ep)
+
+        self.session = InferenceSession(config['model_path'],
+                                        sess_options=sess_opt,
+                                        providers=EP_list)
+
+        if config['use_cuda'] and cuda_ep not in self.session.get_providers():
+            warnings.warn(f'{cuda_ep} is not avaiable for current env, the inference part is automatically shifted to be executed under {cpu_ep}.\n'
+                          'Please ensure the installed onnxruntime-gpu version matches your cuda and cudnn version, '
+                          'you can check their relations from the offical web site: '
+                          '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 get_output_name(self, output_idx=0):
+        return self.session.get_outputs()[output_idx].name
 
 
 def read_yaml(yaml_path):

python/rapidocr_openvino/ch_ppocr_v2_det/text_detect.py

@@ -17,14 +17,13 @@ import time
 
 import cv2
 import numpy as np
-from openvino.runtime import Core
 
 try:
     from .utils import (DBPostProcess, create_operators,
-                        draw_text_det_res, transform, read_yaml)
+                        transform, OpenVINOInferSession)
 except:
     from utils import (DBPostProcess, create_operators,
-                       draw_text_det_res, transform, read_yaml)
+                       transform, OpenVINOInferSession)
 
 
 class TextDetector(object):
@@ -32,10 +31,29 @@ class TextDetector(object):
         self.preprocess_op = create_operators(config['pre_process'])
         self.postprocess_op = DBPostProcess(**config['post_process'])
 
-        ie = Core()
-        model_onnx = ie.read_model(config['model_path'])
-        compile_model = ie.compile_model(model=model_onnx, device_name='CPU')
-        self.session = compile_model.create_infer_request()
+        openvino_instance = OpenVINOInferSession(config)
+        self.session = openvino_instance.session
+
+    def __call__(self, img):
+        ori_im = img.copy()
+        data = {'image': img}
+        data = transform(data, self.preprocess_op)
+        img, shape_list = data
+        if img is None:
+            return None, 0
+
+        img = np.expand_dims(img, axis=0).astype(np.float32)
+        shape_list = np.expand_dims(shape_list, axis=0)
+
+        starttime = time.time()
+        self.session.infer(inputs=[img])
+        preds = self.session.get_output_tensor().data
+
+        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)
+        elapse = time.time() - starttime
+        return dt_boxes, elapse
 
     def order_points_clockwise(self, pts):
         """
@@ -81,27 +99,6 @@ class TextDetector(object):
         dt_boxes = np.array(dt_boxes_new)
         return dt_boxes
 
-    def __call__(self, img):
-        ori_im = img.copy()
-        data = {'image': img}
-        data = transform(data, self.preprocess_op)
-        img, shape_list = data
-        if img is None:
-            return None, 0
-
-        img = np.expand_dims(img, axis=0).astype(np.float32)
-        shape_list = np.expand_dims(shape_list, axis=0)
-
-        starttime = time.time()
-        self.session.infer(inputs=[img])
-        preds = self.session.get_output_tensor().data
-
-        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)
-        elapse = time.time() - starttime
-        return dt_boxes, elapse
-
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -109,12 +106,15 @@ if __name__ == "__main__":
     parser.add_argument('--image_path', type=str, default=None)
     args = parser.parse_args()
 
+    from utils import read_yaml
     config = read_yaml(args.config_path)
 
     text_detector = TextDetector(config)
 
     img = cv2.imread(args.image_path)
     dt_boxes, elapse = text_detector(img)
+
+    from utils import draw_text_det_res
     src_im = draw_text_det_res(dt_boxes, args.image_path)
     cv2.imwrite('det_results.jpg', src_im)
     print('图像已经保存为det_results.jpg了')

python/rapidocr_openvino/ch_ppocr_v2_det/utils.py

@@ -13,11 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-# coding:utf-8
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
 from __future__ import (absolute_import, division, print_function,
                         unicode_literals)
 
-import os
 import sys
 
 import cv2
@@ -26,6 +27,15 @@ import pyclipper
 import six
 from shapely.geometry import Polygon
 import yaml
+from openvino.runtime import Core
+
+
+class OpenVINOInferSession(object):
+    def __init__(self, config):
+        ie = Core()
+        model_onnx = ie.read_model(config['model_path'])
+        compile_model = ie.compile_model(model=model_onnx, device_name='CPU')
+        self.session = compile_model.create_infer_request()
 
 
 def read_yaml(yaml_path):

python/rapidocr_openvino/ch_ppocr_v2_rec/text_recognize.py

@@ -11,20 +11,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
 import argparse
 import math
 import time
-from pathlib import Path
 from typing import List
 
 import cv2
 import numpy as np
-from openvino.runtime import Core
 
 try:
-    from .utils import CTCLabelDecode, read_yaml
+    from .utils import CTCLabelDecode, OpenVINOInferSession
 except:
-    from utils import CTCLabelDecode, read_yaml
+    from utils import CTCLabelDecode, OpenVINOInferSession
 
 
 class TextRecognizer(object):
@@ -34,35 +35,8 @@ class TextRecognizer(object):
         self.character_dict_path =  config['keys_path']
         self.postprocess_op = CTCLabelDecode(self.character_dict_path)
 
-        ie = Core()
-        model_onnx = ie.read_model(config['model_path'])
-        compile_model = ie.compile_model(model=model_onnx, device_name='CPU')
-        self.session = compile_model.create_infer_request()
-
-    def resize_norm_img(self, img, max_wh_ratio):
-        img_channel, img_height, img_width = self.rec_image_shape
-        assert img_channel == img.shape[2]
-
-        img_width = int((32 * max_wh_ratio))
-        max_wh_ratio = 1
-
-        h, w = img.shape[:2]
-        ratio = w / float(h)
-        if math.ceil(img_height * ratio) > img_width:
-            resized_w = img_width
-        else:
-            resized_w = int(math.ceil(img_height * ratio))
-
-        resized_image = cv2.resize(img, (resized_w, img_height))
-        resized_image = resized_image.astype('float32')
-        resized_image = resized_image.transpose((2, 0, 1)) / 255
-        resized_image -= 0.5
-        resized_image /= 0.5
-
-        padding_im = np.zeros((img_channel, img_height, img_width),
-                              dtype=np.float32)
-        padding_im[:, :, 0:resized_w] = resized_image
-        return padding_im
+        openvino_instance = OpenVINOInferSession(config)
+        self.session = openvino_instance.session
 
     def __call__(self, img_list: List[np.ndarray]):
         if isinstance(img_list, np.ndarray):
@@ -104,6 +78,31 @@ class TextRecognizer(object):
             elapse += time.time() - starttime
         return rec_res, elapse
 
+    def resize_norm_img(self, img, max_wh_ratio):
+        img_channel, img_height, img_width = self.rec_image_shape
+        assert img_channel == img.shape[2]
+
+        img_width = int((img_height * max_wh_ratio))
+        max_wh_ratio = 1
+
+        h, w = img.shape[:2]
+        ratio = w / float(h)
+        if math.ceil(img_height * ratio) > img_width:
+            resized_w = img_width
+        else:
+            resized_w = int(math.ceil(img_height * ratio))
+
+        resized_image = cv2.resize(img, (resized_w, img_height))
+        resized_image = resized_image.astype('float32')
+        resized_image = resized_image.transpose((2, 0, 1)) / 255
+        resized_image -= 0.5
+        resized_image /= 0.5
+
+        padding_im = np.zeros((img_channel, img_height, img_width),
+                              dtype=np.float32)
+        padding_im[:, :, 0:resized_w] = resized_image
+        return padding_im
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -111,9 +110,10 @@ if __name__ == "__main__":
     parser.add_argument('--config_path', type=str, default='config.yaml')
     args = parser.parse_args()
 
+    from utils import read_yaml
     config = read_yaml(args.config_path)
     text_recognizer = TextRecognizer(config)
 
     img = cv2.imread(args.image_path)
     rec_res, predict_time = text_recognizer(img)
-    print(f'识别结果: {rec_res}\t cost: {predict_time}s')
+    print(f'rec result: {rec_res}\t cost: {predict_time}s')

python/rapidocr_openvino/ch_ppocr_v2_rec/utils.py

@@ -3,6 +3,15 @@
 # @Contact: liekkaskono@163.com
 import numpy as np
 import yaml
+from openvino.runtime import Core
+
+
+class OpenVINOInferSession(object):
+    def __init__(self, config):
+        ie = Core()
+        model_onnx = ie.read_model(config['model_path'])
+        compile_model = ie.compile_model(model=model_onnx, device_name='CPU')
+        self.session = compile_model.create_infer_request()
 
 
 def read_yaml(yaml_path):

python/rapidocr_openvino/rapid_ocr_api.py

@@ -14,8 +14,10 @@ sys.path.append(str(root_dir))
 
 
 class TextSystem(object):
-    def __init__(self, config_path=str(root_dir / 'config.yaml')):
+    def __init__(self, config_path):
         super(TextSystem).__init__()
+        if not Path(config_path).exists():
+            raise FileExistsError(f'{config_path} does not exist!')
         config = self.read_yaml(config_path)
 
         global_config = config['Global']