From c1fa81f1765e28eb33d36cb97f2eec256067e6f2 Mon Sep 17 00:00:00 2001 From: SWHL Date: Wed, 23 Nov 2022 14:21:31 +0800 Subject: [PATCH] Add verify model function. --- .../ch_ppocr_v2_cls/utils.py | 10 ++++ .../ch_ppocr_v3_det/utils.py | 10 ++++ .../ch_ppocr_v3_rec/text_recognize.py | 53 +++++++++---------- .../ch_ppocr_v3_rec/utils.py | 21 ++++++-- .../ch_ppocr_v2_cls/utils.py | 10 ++++ .../ch_ppocr_v3_det/utils.py | 10 ++++ .../ch_ppocr_v3_rec/text_recognize.py | 5 +- .../ch_ppocr_v3_rec/utils.py | 21 ++++++-- .../ch_ppocr_v2_cls/utils.py | 11 ++++ .../ch_ppocr_v3_det/utils.py | 16 ++++-- .../ch_ppocr_v3_rec/text_recognize.py | 2 +- .../ch_ppocr_v3_rec/utils.py | 10 ++++ 12 files changed, 138 insertions(+), 41 deletions(-) diff --git a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py index 93f09ba..a15cd65 100644 --- a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py +++ b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py @@ -11,6 +11,7 @@ # 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. +from pathlib import Path import warnings import yaml @@ -38,6 +39,7 @@ class OrtInferSession(object): EP_list = [(cuda_ep, config[cuda_ep])] EP_list.append((cpu_ep, cpu_provider_options)) + self._verify_model(config['model_path']) self.session = InferenceSession(config['model_path'], sess_options=sess_opt, providers=EP_list) @@ -55,6 +57,14 @@ class OrtInferSession(object): def get_output_name(self, output_idx=0): return self.session.get_outputs()[output_idx].name + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: diff --git a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py index 75e7bb7..c2c9e17 100644 --- a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py +++ b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py @@ -16,6 +16,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from pathlib import Path import sys import warnings @@ -49,6 +50,7 @@ class OrtInferSession(object): EP_list = [(cuda_ep, config[cuda_ep])] EP_list.append((cpu_ep, cpu_provider_options)) + self._verify_model(config['model_path']) self.session = InferenceSession(config['model_path'], sess_options=sess_opt, providers=EP_list) @@ -66,6 +68,14 @@ class OrtInferSession(object): def get_output_name(self, output_idx=0): return self.session.get_outputs()[output_idx].name + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: diff --git a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py index 753c811..4d306aa 100644 --- a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py +++ b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py @@ -30,10 +30,9 @@ class TextRecognizer(object): session_instance = OrtInferSession(config) self.session = session_instance.session self.input_name = session_instance.get_input_name() - meta_dict = session_instance.get_metadata() - if 'character' in meta_dict.keys(): - self.character_dict_path = meta_dict['character'].splitlines() + if session_instance.have_key(): + self.character_dict_path = session_instance.get_character_list() else: self.character_dict_path = config.get('keys_path', None) self.postprocess_op = CTCLabelDecode(self.character_dict_path) @@ -41,30 +40,6 @@ class TextRecognizer(object): self.rec_batch_num = config['rec_batch_num'] self.rec_image_shape = config['rec_img_shape'] - def resize_norm_img(self, img, max_wh_ratio): - img_channel, img_height, img_width = self.rec_image_shape - assert img_channel == img.shape[2] - - img_width = int(img_height * max_wh_ratio) - - h, w = img.shape[:2] - ratio = w / float(h) - if math.ceil(img_height * ratio) > img_width: - resized_w = img_width - else: - resized_w = int(math.ceil(img_height * ratio)) - - resized_image = cv2.resize(img, (resized_w, img_height)) - resized_image = resized_image.astype('float32') - resized_image = resized_image.transpose((2, 0, 1)) / 255 - resized_image -= 0.5 - resized_image /= 0.5 - - padding_im = np.zeros((img_channel, img_height, img_width), - dtype=np.float32) - padding_im[:, :, 0: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] @@ -105,6 +80,30 @@ class TextRecognizer(object): elapse += time.time() - starttime return rec_res, elapse + def resize_norm_img(self, img, max_wh_ratio): + img_channel, img_height, img_width = self.rec_image_shape + assert img_channel == img.shape[2] + + img_width = int(img_height * max_wh_ratio) + + h, w = img.shape[:2] + ratio = w / float(h) + if math.ceil(img_height * ratio) > img_width: + resized_w = img_width + else: + resized_w = int(math.ceil(img_height * ratio)) + + resized_image = cv2.resize(img, (resized_w, img_height)) + resized_image = resized_image.astype('float32') + resized_image = resized_image.transpose((2, 0, 1)) / 255 + resized_image -= 0.5 + resized_image /= 0.5 + + padding_im = np.zeros((img_channel, img_height, img_width), + dtype=np.float32) + padding_im[:, :, 0:resized_w] = resized_image + return padding_im + if __name__ == "__main__": parser = argparse.ArgumentParser() diff --git a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py index 6aecf86..0a11e42 100644 --- a/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py +++ b/ocrweb/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from pathlib import Path import warnings import numpy as np @@ -29,6 +30,7 @@ class OrtInferSession(object): EP_list = [(cuda_ep, config[cuda_ep])] EP_list.append((cpu_ep, cpu_provider_options)) + self._verify_model(config['model_path']) self.session = InferenceSession(config['model_path'], sess_options=sess_opt, providers=EP_list) @@ -46,9 +48,22 @@ class OrtInferSession(object): def get_output_name(self, output_idx=0): return self.session.get_outputs()[output_idx].name - def get_metadata(self): - meta_dict = self.session.get_modelmeta().custom_metadata_map - return meta_dict + def get_character_list(self, key: str = 'character'): + return self.meta_dict[key].splitlines() + + def have_key(self, key: str = 'character') -> bool: + self.meta_dict = self.session.get_modelmeta().custom_metadata_map + if key in self.meta_dict.keys(): + return True + return False + + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') def read_yaml(yaml_path): diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py index 93f09ba..a15cd65 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_cls/utils.py @@ -11,6 +11,7 @@ # 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. +from pathlib import Path import warnings import yaml @@ -38,6 +39,7 @@ class OrtInferSession(object): EP_list = [(cuda_ep, config[cuda_ep])] EP_list.append((cpu_ep, cpu_provider_options)) + self._verify_model(config['model_path']) self.session = InferenceSession(config['model_path'], sess_options=sess_opt, providers=EP_list) @@ -55,6 +57,14 @@ class OrtInferSession(object): def get_output_name(self, output_idx=0): return self.session.get_outputs()[output_idx].name + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py index 75e7bb7..c2c9e17 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_det/utils.py @@ -16,6 +16,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from pathlib import Path import sys import warnings @@ -49,6 +50,7 @@ class OrtInferSession(object): EP_list = [(cuda_ep, config[cuda_ep])] EP_list.append((cpu_ep, cpu_provider_options)) + self._verify_model(config['model_path']) self.session = InferenceSession(config['model_path'], sess_options=sess_opt, providers=EP_list) @@ -66,6 +68,14 @@ class OrtInferSession(object): def get_output_name(self, output_idx=0): return self.session.get_outputs()[output_idx].name + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: 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 f05fe30..df1bda0 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py @@ -30,10 +30,9 @@ class TextRecognizer(object): session_instance = OrtInferSession(config) self.session = session_instance.session self.input_name = session_instance.get_input_name() - meta_dict = session_instance.get_metadata() - if 'character' in meta_dict.keys(): - self.character_dict_path = meta_dict['character'].splitlines() + if session_instance.have_key(): + self.character_dict_path = session_instance.get_character_list() else: self.character_dict_path = config.get('keys_path', None) self.postprocess_op = CTCLabelDecode(self.character_dict_path) diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py index 6aecf86..0a11e42 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from pathlib import Path import warnings import numpy as np @@ -29,6 +30,7 @@ class OrtInferSession(object): EP_list = [(cuda_ep, config[cuda_ep])] EP_list.append((cpu_ep, cpu_provider_options)) + self._verify_model(config['model_path']) self.session = InferenceSession(config['model_path'], sess_options=sess_opt, providers=EP_list) @@ -46,9 +48,22 @@ class OrtInferSession(object): def get_output_name(self, output_idx=0): return self.session.get_outputs()[output_idx].name - def get_metadata(self): - meta_dict = self.session.get_modelmeta().custom_metadata_map - return meta_dict + def get_character_list(self, key: str = 'character'): + return self.meta_dict[key].splitlines() + + def have_key(self, key: str = 'character') -> bool: + self.meta_dict = self.session.get_modelmeta().custom_metadata_map + if key in self.meta_dict.keys(): + return True + return False + + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') def read_yaml(yaml_path): diff --git a/python/rapidocr_openvino/ch_ppocr_v2_cls/utils.py b/python/rapidocr_openvino/ch_ppocr_v2_cls/utils.py index abcf3c3..37ce4c4 100644 --- a/python/rapidocr_openvino/ch_ppocr_v2_cls/utils.py +++ b/python/rapidocr_openvino/ch_ppocr_v2_cls/utils.py @@ -11,6 +11,7 @@ # 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. +from pathlib import Path import yaml from openvino.runtime import Core @@ -18,10 +19,20 @@ from openvino.runtime import Core class OpenVINOInferSession(object): def __init__(self, config): ie = Core() + + self._verify_model(config['model_path']) model_onnx = ie.read_model(config['model_path']) compile_model = ie.compile_model(model=model_onnx, device_name='CPU') self.session = compile_model.create_infer_request() + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: diff --git a/python/rapidocr_openvino/ch_ppocr_v3_det/utils.py b/python/rapidocr_openvino/ch_ppocr_v3_det/utils.py index 2028a4f..3d57926 100644 --- a/python/rapidocr_openvino/ch_ppocr_v3_det/utils.py +++ b/python/rapidocr_openvino/ch_ppocr_v3_det/utils.py @@ -16,27 +16,35 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com -from __future__ import (absolute_import, division, print_function, - unicode_literals) - import sys +from pathlib import Path import cv2 import numpy as np import pyclipper import six -from shapely.geometry import Polygon import yaml from openvino.runtime import Core +from shapely.geometry import Polygon class OpenVINOInferSession(object): def __init__(self, config): ie = Core() + + self._verify_model(config['model_path']) model_onnx = ie.read_model(config['model_path']) compile_model = ie.compile_model(model=model_onnx, device_name='CPU') self.session = compile_model.create_infer_request() + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: diff --git a/python/rapidocr_openvino/ch_ppocr_v3_rec/text_recognize.py b/python/rapidocr_openvino/ch_ppocr_v3_rec/text_recognize.py index 06c4831..ece6ea5 100644 --- a/python/rapidocr_openvino/ch_ppocr_v3_rec/text_recognize.py +++ b/python/rapidocr_openvino/ch_ppocr_v3_rec/text_recognize.py @@ -34,7 +34,7 @@ class TextRecognizer(object): self.rec_image_shape = config['rec_img_shape'] self.rec_batch_num = config['rec_batch_num'] - dict_path = str(Path(__file__).parent / './ppocr_keys_v1.txt') + dict_path = str(Path(__file__).parent / 'ppocr_keys_v1.txt') self.character_dict_path = config.get( 'keys_path', dict_path) self.postprocess_op = CTCLabelDecode(self.character_dict_path) diff --git a/python/rapidocr_openvino/ch_ppocr_v3_rec/utils.py b/python/rapidocr_openvino/ch_ppocr_v3_rec/utils.py index b5cd0eb..c53c224 100644 --- a/python/rapidocr_openvino/ch_ppocr_v3_rec/utils.py +++ b/python/rapidocr_openvino/ch_ppocr_v3_rec/utils.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +from pathlib import Path import numpy as np import yaml from openvino.runtime import Core @@ -9,10 +10,19 @@ from openvino.runtime import Core class OpenVINOInferSession(object): def __init__(self, config): ie = Core() + self._verify_model(config['model_path']) model_onnx = ie.read_model(config['model_path']) compile_model = ie.compile_model(model=model_onnx, device_name='CPU') self.session = compile_model.create_infer_request() + @staticmethod + def _verify_model(model_path): + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f'{model_path} does not exists.') + if not model_path.is_file(): + raise FileExistsError(f'{model_path} is not a file.') + def read_yaml(yaml_path): with open(yaml_path, 'rb') as f: