diff --git a/python/config.yaml b/python/config.yaml index 0c59831..6d177e7 100644 --- a/python/config.yaml +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/config.yaml index e6fa00e..89497ae 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/config.yaml +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py index fc06ac1..38d8549 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py +++ b/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]}") diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py index dae9175..55b4816 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py +++ b/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): diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml index 0390513..428ccfc 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py index d656616..7308253 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py +++ b/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.') diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py index b1fa89d..e743908 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py +++ b/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): diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml index 3c3c9df..ffcfc1d 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py index 2753ea7..4fd9265 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/utils.py index c6edeb6..489a04d 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/utils.py +++ b/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): diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/word_1.jpg b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/word_1.jpg deleted file mode 100644 index 934cc1c..0000000 Binary files a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/word_1.jpg and /dev/null differ diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml index 4f903d6..2428429 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py index ca88ce5..b00bf95 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py index c6edeb6..e50bc2a 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py +++ b/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): diff --git a/python/test_demo.py b/python/test_demo.py index a7697ad..67c10f8 100644 --- a/python/test_demo.py +++ b/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)