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 3d2b673..a46b77e 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py @@ -33,9 +33,7 @@ class TextClassifier(): self.cls_thresh = config['cls_thresh'] self.postprocess_op = ClsPostProcess(config['label_list']) - session_instance = OrtInferSession(config) - self.session = session_instance.session - self.input_name = session_instance.get_input_name() + self.infer = OrtInferSession(config) def __call__(self, img_list: List[np.ndarray]): if isinstance(img_list, np.ndarray): @@ -64,8 +62,7 @@ class TextClassifier(): norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32) starttime = time.time() - onnx_inputs = {self.input_name: norm_img_batch} - prob_out = self.session.run(None, onnx_inputs)[0] + prob_out = self.infer(norm_img_batch)[0] cls_result = self.postprocess_op(prob_out) elapse += time.time() - starttime diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py index d37221e..d2df500 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py @@ -15,6 +15,7 @@ from pathlib import Path import warnings import yaml +import numpy as np from onnxruntime import (get_available_providers, get_device, SessionOptions, InferenceSession, GraphOptimizationLevel) @@ -53,11 +54,18 @@ class OrtInferSession(): '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 __call__(self, input_content: np.ndarray) -> np.ndarray: + input_dict = dict(zip(self.get_input_names(), input_content)) + try: + return self.session.run(self.get_output_names(), input_dict) + except Exception as e: + raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e - def get_output_name(self, output_idx=0): - return self.session.get_outputs()[output_idx].name + def get_input_names(self, ): + return [v.name for v in self.session.get_inputs()] + + def get_output_names(self,): + return [v.name for v in self.session.get_outputs()] @staticmethod def _verify_model(model_path): @@ -68,6 +76,10 @@ class OrtInferSession(): raise FileExistsError(f'{model_path} is not a file.') +class ONNXRuntimeError(Exception): + pass + + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: data = yaml.load(f, Loader=yaml.Loader) diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py index c3ebf6e..38cd186 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_det/text_detect.py @@ -33,11 +33,9 @@ class TextDetector(): self.preprocess_op = create_operators(config['pre_process']) self.postprocess_op = DBPostProcess(**config['post_process']) - session_instance = OrtInferSession(config) - self.session = session_instance.session - self.input_name = session_instance.get_input_name() + self.infer = OrtInferSession(config) - def __call__(self, img): + def __call__(self, img: np.ndarry): if img is None: raise ValueError('img is None') @@ -53,9 +51,8 @@ class TextDetector(): shape_list = np.expand_dims(shape_list, axis=0) starttime = time.time() - preds = self.session.run(None, {self.input_name: img}) - - post_result = self.postprocess_op(preds[0], shape_list) + preds = self.infer(img)[0] + 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) diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py index dc7eccd..110c14d 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py @@ -64,11 +64,18 @@ class OrtInferSession(): '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 __call__(self, input_content: np.ndarray) -> np.ndarray: + input_dict = dict(zip(self.get_input_names(), input_content)) + try: + return self.session.run(self.get_output_names(), input_dict) + except Exception as e: + raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e - def get_output_name(self, output_idx=0): - return self.session.get_outputs()[output_idx].name + def get_input_names(self, ) + return [v.name for v in self.session.get_inputs()] + + def get_output_names(self,): + return [v.name for v in self.session.get_outputs()] @staticmethod def _verify_model(model_path): @@ -79,6 +86,10 @@ class OrtInferSession(): raise FileExistsError(f'{model_path} is not a file.') +class ONNXRuntimeError(Exception): + pass + + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: data = yaml.load(f, Loader=yaml.Loader) 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 116eabb..e996420 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py @@ -27,12 +27,10 @@ except: class TextRecognizer(): def __init__(self, config): - session_instance = OrtInferSession(config) - self.session = session_instance.session - self.input_name = session_instance.get_input_name() + self.session = OrtInferSession(config) - if session_instance.have_key(): - self.character_dict_path = session_instance.get_character_list() + if self.session.have_key(): + self.character_dict_path = self.session.get_character_list() else: self.character_dict_path = config.get('keys_path', None) self.postprocess_op = CTCLabelDecode(self.character_dict_path) @@ -71,8 +69,7 @@ class TextRecognizer(): norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32) starttime = time.time() - onnx_inputs = {self.input_name: norm_img_batch} - preds = self.session.run(None, onnx_inputs)[0] + preds = self.session(norm_img_batch)[0] rec_result = self.postprocess_op(preds) for rno in range(len(rec_result)): diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py index 8d04f25..9198df4 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py @@ -44,11 +44,18 @@ class OrtInferSession(): '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 __call__(self, input_content: np.ndarray) -> np.ndarray: + input_dict = dict(zip(self.get_input_names(), input_content)) + try: + return self.session.run(self.get_output_names(), input_dict) + except Exception as e: + raise ONNXRuntimeError('ONNXRuntime inferece failed.') from e - def get_output_name(self, output_idx=0): - return self.session.get_outputs()[output_idx].name + def get_input_names(self, ): + return [v.name for v in self.session.get_inputs()] + + def get_output_names(self,): + return [v.name for v in self.session.get_outputs()] def get_character_list(self, key: str = 'character'): return self.meta_dict[key].splitlines() @@ -68,6 +75,10 @@ class OrtInferSession(): raise FileExistsError(f'{model_path} is not a file.') +class ONNXRuntimeError(Exception): + pass + + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: data = yaml.load(f, Loader=yaml.Loader) diff --git a/python/tests/test_all_ort.py b/python/tests/test_all_ort.py index 1a77900..be68e7a 100644 --- a/python/tests/test_all_ort.py +++ b/python/tests/test_all_ort.py @@ -15,7 +15,7 @@ sys.path.append(str(root_dir)) from rapidocr_onnxruntime import RapidOCR rapid_ocr = RapidOCR() -tests_dir = root_dir / 'test_files' +tests_dir = root_dir / 'tests' / 'test_files' def test_normal(): diff --git a/python/tests/test_all_vino.py b/python/tests/test_all_vino.py index fe1cca0..479fea0 100644 --- a/python/tests/test_all_vino.py +++ b/python/tests/test_all_vino.py @@ -16,7 +16,7 @@ from rapidocr_openvino import RapidOCR rapid_ocr = RapidOCR() -tests_dir = root_dir / 'test_files' +tests_dir = root_dir / 'tests' / 'test_files' def test_normal():