Add use_cuda parameter

python/config.yaml

@@ -8,6 +8,15 @@ Det:
     class_name: TextDetector
     model_path: resources/models/ch_PP-OCRv2_det_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
+        gpu_mem_limit: 2 * 1024 * 1024 * 1024
+        cudnn_conv_algo_search: EXHAUSTIVE
+        do_copy_in_default_stream: true
+
     pre_process:
         DetResizeForTest:
             limit_side_len: 736
@@ -33,6 +42,14 @@ Cls:
     class_name: TextClassifier
     model_path: resources/models/ch_ppocr_mobile_v2.0_cls_infer.onnx
 
+    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
+
     cls_image_shape: [3, 48, 192]
     cls_batch_num: 6
     cls_thresh: 0.9
@@ -43,6 +60,14 @@ Rec:
     class_name: TextRecognizer
     model_path: resources/models/ch_ppocr_mobile_v2.0_rec_infer.onnx
 
+    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
+
     rec_img_shape: [3, 32, 320]
     rec_batch_num: 6
     keys_path: resources/rec_dict/ppocr_keys_v1.txt

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/config.yaml

@@ -1,5 +1,14 @@
 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
+    gpu_mem_limit: 2 * 1024 * 1024 * 1024
+    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_onnxruntime/ch_ppocr_v2_cls/text_cls.py

@@ -22,9 +22,9 @@ 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 +34,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 resize_norm_img(self, img):
         img_c, img_h, img_w = self.cls_image_shape
@@ -63,7 +61,6 @@ class TextClassifier(object):
         padding_im[:, :, :resized_w] = resized_image
         return padding_im
 
-
     def __call__(self, img_list: List[np.ndarray]):
         if isinstance(img_list, np.ndarray):
             img_list = [img_list]
@@ -123,4 +120,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_onnxruntime/ch_ppocr_v2_cls/utils.py

@@ -11,7 +11,47 @@
 # 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])]
+        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)
+
+        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_onnxruntime/ch_ppocr_v2_det/config.yaml

@@ -1,7 +1,7 @@
 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
+# Details of the params: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html
 CUDAExecutionProvider:
     device_id: 0
     arena_extend_strategy: kNextPowerOfTwo

python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py

@@ -37,7 +37,6 @@ class TextDetector(object):
         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):
         """
@@ -100,7 +99,7 @@ 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(None, {self.input_name: img})
 
         post_result = self.postprocess_op(preds[0], shape_list)
 
@@ -124,4 +123,4 @@ if __name__ == "__main__":
     dt_boxes, elapse = text_detector(img)
     src_im = draw_text_det_res(dt_boxes, args.image_path)
     cv2.imwrite('det_results.jpg', src_im)
-    print('图像已经保存为det_results.jpg了')
+    print('The det_results.jpg has been saved in the current directory.')

python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py

@@ -50,10 +50,12 @@ class OrtInferSession(object):
         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',
+            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):

python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml

@@ -1,5 +1,14 @@
 model_path: resources/models/ch_ppocr_server_v2.0_rec_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
+    gpu_mem_limit: 2 * 1024 * 1024 * 1024
+    cudnn_conv_algo_search: EXHAUSTIVE
+    do_copy_in_default_stream: true
+
 rec_img_shape: [3, 32, 320]
 rec_batch_num: 6
 keys_path: resources/rec_dict/ppocr_keys_v1.txt

python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py

@@ -14,17 +14,15 @@
 import argparse
 import math
 import time
-from pathlib import Path
 from typing import List
 
 import cv2
 import numpy as np
-import onnxruntime
 
 try:
-    from .utils import CTCLabelDecode, read_yaml
+    from .utils import CTCLabelDecode, read_yaml, OrtInferSession
 except:
-    from utils import CTCLabelDecode, read_yaml
+    from utils import CTCLabelDecode, read_yaml, OrtInferSession
 
 
 class TextRecognizer(object):
@@ -34,12 +32,9 @@ class TextRecognizer(object):
         self.character_dict_path =  config['keys_path']
         self.postprocess_op = CTCLabelDecode(self.character_dict_path)
 
-        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
+        session_instance = OrtInferSession(config)
+        self.session = session_instance.session
+        self.input_name = session_instance.get_input_name()
 
     def resize_norm_img(self, img, max_wh_ratio):
         img_channel, img_height, img_width = self.rec_image_shape

python/rapidocr_onnxruntime/ch_ppocr_v2_rec/utils.py

@@ -1,8 +1,48 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+import warnings
+
 import numpy as np
 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])]
+        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)
+
+        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_onnxruntime/ch_ppocr_v3_rec/config.yaml

@@ -1,5 +1,14 @@
 model_path: resources/models/PPOCRv3/ch_PP-OCRv3_rec_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
+    gpu_mem_limit: 2 * 1024 * 1024 * 1024
+    cudnn_conv_algo_search: EXHAUSTIVE
+    do_copy_in_default_stream: true
+
 rec_img_shape: [3, 48, 320]
 rec_batch_num: 6
 keys_path: resources/rec_dict/ppocr_keys_v1.txt

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py

@@ -22,9 +22,9 @@ import numpy as np
 import onnxruntime
 
 try:
-    from .utils import CTCLabelDecode, read_yaml
+    from .utils import CTCLabelDecode, read_yaml, OrtInferSession
 except:
-    from utils import CTCLabelDecode, read_yaml
+    from utils import CTCLabelDecode, read_yaml, OrtInferSession
 
 root_path = Path(__file__).parent
 
@@ -36,12 +36,9 @@ class TextRecognizer(object):
         self.character_dict_path = config['keys_path']
         self.postprocess_op = CTCLabelDecode(self.character_dict_path)
 
-        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
+        session_instance = OrtInferSession(config)
+        self.session = session_instance.session
+        self.input_name = session_instance.get_input_name()
 
     def resize_norm_img(self, img, max_wh_ratio):
         img_channel, img_height, img_width = self.rec_image_shape

python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py

@@ -1,8 +1,48 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
+import warnings
+
 import numpy as np
 import yaml
+from onnxruntime import (InferenceSession, SessionOptions,
+                         get_available_providers, get_device)
+
+
+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)
+
+        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/test_demo.py

@@ -89,8 +89,7 @@ def visualize(image_path, boxes, rec_res, font_path="fonts/msyh.ttc"):
 
 
 if __name__ == '__main__':
-    config_path = 'config.yaml'
-    text_sys = TextSystem(config_path)
+    text_sys = TextSystem('config.yaml')
 
     image_path = r'test_images/det_images/ch_en_num.jpg'
     img = cv2.imread(image_path)