The det part almost

python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml

@@ -1,4 +1,13 @@
-model_path: resources/models/ch_ppocr_server_v2.0_det_infer.onnx
+model_path: G:\ProgramFiles\_self\RapidOCR\python\resources\models\ch_PP-OCRv2_det_infer.onnx
+
+# https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html
+use_cuda: false
+CUDAExecutionProvider:
+    device_id: 0
+    arena_extend_strategy: kNextPowerOfTwo
+    gpu_mem_limit: 2 * 1024 * 1024 * 1024
+    cudnn_conv_algo_search: EXHAUSTIVE
+    do_copy_in_default_stream: true
 
 pre_process:
     DetResizeForTest:

python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py

@@ -12,19 +12,21 @@
 # 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 time
 
 import cv2
 import numpy as np
-import onnxruntime
+
 
 try:
     from .utils import (DBPostProcess, create_operators,
-                        draw_text_det_res, transform, read_yaml)
+                        draw_text_det_res, transform, read_yaml, OrtInferSession)
 except:
     from utils import (DBPostProcess, create_operators,
-                       draw_text_det_res, transform, read_yaml)
+                       draw_text_det_res, transform, read_yaml, OrtInferSession)
 
 
 class TextDetector(object):
@@ -32,18 +34,16 @@ class TextDetector(object):
         self.preprocess_op = create_operators(config['pre_process'])
         self.postprocess_op = DBPostProcess(**config['post_process'])
 
-        sess_opt = onnxruntime.SessionOptions()
-        sess_opt.log_severity_level = 4
-        sess_opt.enable_cpu_mem_arena = False
-        self.session = onnxruntime.InferenceSession(config['model_path'],
-                                                    sess_opt)
-        self.input_name = self.session.get_inputs()[0].name
-        self.output_name = self.session.get_outputs()[0].name
+        session_instance = OrtInferSession(config)
+        self.session = session_instance.session
+        self.input_name = session_instance.get_input_name()
+        self.output_name = session_instance.get_output_name()
 
     def order_points_clockwise(self, pts):
         """
-        reference from: https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
-        # sort the points based on their x-coordinates
+        reference from:
+        https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
+        sort the points based on their x-coordinates
         """
         xSorted = pts[np.argsort(pts[:, 0]), :]
 
@@ -85,7 +85,11 @@ class TextDetector(object):
         return dt_boxes
 
     def __call__(self, img):
-        ori_im = img.copy()
+        if img is None:
+            raise ValueError('img is None')
+
+        ori_im_shape = img.shape[:2]
+
         data = {'image': img}
         data = transform(data, self.preprocess_op)
         img, shape_list = data
@@ -96,13 +100,12 @@ class TextDetector(object):
         shape_list = np.expand_dims(shape_list, axis=0)
 
         starttime = time.time()
-        preds = self.session.run([self.output_name],
-                                 {self.input_name: img})
+        preds = self.session.run([self.output_name], {self.input_name: img})
 
         post_result = self.postprocess_op(preds[0], shape_list)
 
         dt_boxes = post_result[0]['points']
-        dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
+        dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im_shape)
         elapse = time.time() - starttime
         return dt_boxes, elapse
 

python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py

@@ -13,19 +13,54 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-# coding:utf-8
-from __future__ import (absolute_import, division, print_function,
-                        unicode_literals)
-
-import os
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
 import sys
+import warnings
 
 import cv2
 import numpy as np
 import pyclipper
 import six
-from shapely.geometry import Polygon
 import yaml
+from shapely.geometry import Polygon
+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])]
+        else:
+            raise RuntimeError(f'{cuda_ep} is not avaiable for --use_cuda!')
+
+        EP_list.append(cpu_ep)
+
+        self.session = InferenceSession(config['model_path'],
+                                        sess_options=sess_opt,
+                                        providers=EP_list)
+        # https://github.com/microsoft/onnxruntime/issues/11323
+        # https://github.com/microsoft/onnxruntime/issues/11996
+        if config['use_cuda'] and cuda_ep not in self.session.get_providers():
+            warnings.warn(f'{cuda_ep} is not avaiable for current env',
+                          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):
@@ -227,6 +262,7 @@ def transform(data, ops=None):
     """ transform """
     if ops is None:
         ops = []
+
     for op in ops:
         data = op(data)
         if data is None: