diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml index 0f08b09..0390513 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml +++ b/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: 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 30af2e3..d656616 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py +++ b/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 diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py index ffbd17f..b1fa89d 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py +++ b/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: