From a3d3f12be7513220f1e8afdf86fd01dcae2a99de Mon Sep 17 00:00:00 2001 From: SWHL Date: Thu, 7 Jul 2022 22:03:51 +0800 Subject: [PATCH] Combine rec v2 and v3 --- python/README.md | 39 +++--- python/config.yaml | 4 +- .../ch_ppocr_v2_rec/__init__.py | 3 + .../ch_ppocr_v2_rec/config.yaml | 4 +- .../ch_ppocr_v2_rec/text_recognize.py | 50 +++---- .../ch_ppocr_v3_rec/README.md | 9 -- .../ch_ppocr_v3_rec/__init__.py | 4 - .../ch_ppocr_v3_rec/config.yaml | 14 -- .../ch_ppocr_v3_rec/text_recognize.py | 130 ------------------ .../ch_ppocr_v3_rec/utils.py | 115 ---------------- python/rapidocr_onnxruntime/rapid_ocr_api.py | 8 +- python/tests/{test_rec_v2.py => test_rec.py} | 0 python/tests/test_rec_v3.py | 24 ---- 13 files changed, 59 insertions(+), 345 deletions(-) delete mode 100644 python/rapidocr_onnxruntime/ch_ppocr_v3_rec/README.md delete mode 100644 python/rapidocr_onnxruntime/ch_ppocr_v3_rec/__init__.py delete mode 100644 python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml delete mode 100644 python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py delete mode 100644 python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py rename python/tests/{test_rec_v2.py => test_rec.py} (100%) delete mode 100644 python/tests/test_rec_v3.py diff --git a/python/README.md b/python/README.md index de47f38..8a052be 100755 --- a/python/README.md +++ b/python/README.md @@ -17,24 +17,37 @@ ### 简介和说明 - **各个版本的ONNX模型下载地址：**[百度网盘](https://pan.baidu.com/s/1PTcgXG2zEgQU6A_A3kGJ3Q?pwd=jhai) | [Google Drive](https://drive.google.com/drive/folders/1x_a9KpCo_1blxH1xFOfgKVkw1HYRVywY?usp=sharing) -- 所有常用的参数配置都在`config.yaml`下，一目了然，更加便捷 +- 所有常用的参数配置都在[`config.yaml`](./config.yaml)下，一目了然，更加便捷 - 每个独立的模块下均有独立的`config.yaml`配置文件，可以单独使用 - `det`部分： - - `det`中`mobile`和`server`版，共用一个推理代码，直接更改配置文件中模型路径即可 - - `det`中`v2`和`v3`两个版本代码，共用一个推理代码，直接更改配置文件中模型路径即可 + - `det`中`mobile`和`server`版，推理代码一致，直接更改配置文件中模型路径即可 + - `det`中`v2`和`v3`两个版本，推理代码一致，直接更改配置文件中模型路径即可 ```yaml Det: module_name: ch_ppocr_v2_det class_name: TextDetector model_path: resources/models/ch_PP-OCRv3_det_infer.onnx ``` -- `rec`中`mobile`和`server`版本，共用一个推理代码，直接更改配置文件中模型路径即可 -- onnxruntime和OpenVINO调用方式如下: +- `rec`部分： + - `rec`中`mobile`和`server`版本，推理代码一致，直接更改配置文件中模型路径即可 + - `rec`中`v2`和`v3`两个版本，共用同一个推理代码。 + - 两版本差别仅在输入shape和模型。经过测试，采用`v3 rec模型`+`[3, 48, 320]`效果最好。 + - 目前配置文件`config.yaml`中（如下所示）已经改为该最优组合。 + ```yaml + module_name: ch_ppocr_v2_rec + class_name: TextRecognizer + model_path: resources/models/ch_PP-OCRv3_rec_infer.onnx + + rec_img_shape: [3, 48, 320] + rec_batch_num: 6 + keys_path: resources/rec_dict/ppocr_keys_v1.txt + ``` +- onnxruntime和openvino调用方式如下: ```python # 基于onnxruntime引擎推理 from rapidocr_onnxruntime import TextSystem - # 基于openvin引擎推理 + # 基于openvino引擎推理 from rapidocr_openvino import TextSystem ``` - 值得说明的是，基于openvino推理部分中`ch_ppocr_v2_cls`部分仍然是基于onnxruntime的，原因是openvino有bug，详情见[openvino/issue](https://github.com/openvinotoolkit/openvino/issues/11501) @@ -92,7 +105,7 @@ pip install -r requirements.txt -i https://pypi.douban.com/simple/ ``` - - 基于OpenVINO推理所需环境安装： + - 基于openvino推理所需环境安装： ```bash # Windows端 pip install openvino==2022.1.0 @@ -114,7 +127,7 @@ # 基于onnxruntime引擎推理 from rapidocr_onnxruntime import TextSystem - # 基于OpenVINO引擎推理 + # 基于openvino引擎推理 # from rapidocr_openvino import TextSystem config_path = 'config.yaml' @@ -129,16 +142,6 @@ ```bash python test_demp.py ``` - - `rec`部分，如果想要使用`PPOCR-v3`的话，需要将`config.yaml`中的`Rec`部分改为如下: - ```yaml - module_name: ch_ppocr_v3_rec - class_name: TextRecognizer - model_path: resources/models/PPOCRv3/ch_PP-OCRv3_rec_infer.onnx - - rec_img_shape: [3, 48, 320] - rec_batch_num: 6 - keys_path: resources/rec_dict/ppocr_keys_v1.txt - ``` ### [`config.yaml`](./config.yaml)中常用参数介绍 - `Global`部分 diff --git a/python/config.yaml b/python/config.yaml index 891d010..a55a535 100644 --- a/python/config.yaml +++ b/python/config.yaml @@ -57,7 +57,7 @@ Cls: Rec: module_name: ch_ppocr_v2_rec class_name: TextRecognizer - model_path: resources/models/ch_ppocr_mobile_v2.0_rec_infer.onnx + model_path: resources/models/ch_PP-OCRv3_rec_infer.onnx use_cuda: false CUDAExecutionProvider: @@ -66,6 +66,6 @@ Rec: cudnn_conv_algo_search: EXHAUSTIVE do_copy_in_default_stream: true - rec_img_shape: [3, 32, 320] + 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_v2_rec/__init__.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/__init__.py index 686ffdd..ea54cd3 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/__init__.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/__init__.py @@ -1 +1,4 @@ +# -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com from .text_recognize import TextRecognizer \ No newline at end of file diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml index 035d74e..c07c65b 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/config.yaml @@ -1,4 +1,4 @@ -model_path: resources/models/ch_ppocr_mobile_v2.0_rec_infer.onnx +model_path: resources/models/ch_PP-OCRv3_rec_infer.onnx use_cuda: false # Details of the params: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html @@ -9,6 +9,6 @@ CUDAExecutionProvider: cudnn_conv_algo_search: EXHAUSTIVE do_copy_in_default_stream: true -rec_img_shape: [3, 32, 320] +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_v2_rec/text_recognize.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py index 4fd9265..b754c4d 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_rec/text_recognize.py @@ -36,31 +36,6 @@ class TextRecognizer(object): 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 - assert img_channel == img.shape[2] - - img_width = int((32 * max_wh_ratio)) - max_wh_ratio = 1 - - 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] @@ -101,6 +76,31 @@ 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)) + max_wh_ratio = 1 + + 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/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/README.md b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/README.md deleted file mode 100644 index cfef298..0000000 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/README.md +++ /dev/null @@ -1,9 +0,0 @@ - -#### 简单对比 - -|输入尺寸|前处理函数|模型|测试集|Score|Exact_Match|Char_Match|Rec_Speed(s/条)| -|---|---|---|---|---|---|---|---| -|48 |resize_norm_img_svtr |ch_PP-OCRv3_rec_infer.onnx|中英印刷体|0.5011|0.3571|0.6450|0.0203| -|48 |resize_norm_img |ch_PP-OCRv3_rec_infer.onnx|中英印刷体|0.7248|0.5417|0.9079|0.0209| -|32 | resize_norm_img_svtr |ch_PP-OCRv3_rec_infer.onnx|中英印刷体|0.3019|0.1607|0.4431|0.0161| -|32| resize_norm_img |ch_PP-OCRv3_rec_infer.onnx|中英印刷体|0.5537|0.3274|0.7800|0.0152| diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/__init__.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/__init__.py deleted file mode 100644 index ea54cd3..0000000 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -# -*- encoding: utf-8 -*- -# @Author: SWHL -# @Contact: liekkaskono@163.com -from .text_recognize import TextRecognizer \ No newline at end of file diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml deleted file mode 100644 index c07c65b..0000000 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/config.yaml +++ /dev/null @@ -1,14 +0,0 @@ -model_path: resources/models/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 deleted file mode 100644 index b00bf95..0000000 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/text_recognize.py +++ /dev/null @@ -1,130 +0,0 @@ -# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# 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 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, OrtInferSession -except: - from utils import CTCLabelDecode, read_yaml, OrtInferSession - -root_path = Path(__file__).parent - - -class TextRecognizer(object): - def __init__(self, config): - self.rec_image_shape = config['rec_img_shape'] - self.rec_batch_num = config['rec_batch_num'] - self.character_dict_path = config['keys_path'] - self.postprocess_op = CTCLabelDecode(self.character_dict_path) - - 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 - assert img_channel == img.shape[2] - - img_width = int((32 * max_wh_ratio)) - max_wh_ratio = 1 - - 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 resize_norm_img_svtr(self, img, max_wh_ratio=None): - img_c, img_h, img_w = self.rec_image_shape - resized_image = cv2.resize(img, (img_w, img_h), - interpolation=cv2.INTER_LINEAR) - - resized_image = resized_image.astype('float32') - resized_image = resized_image.transpose((2, 0, 1)) / 255 - resized_image -= 0.5 - resized_image /= 0.5 - return resized_image - - def __call__(self, img_list: List[np.ndarray]): - if isinstance(img_list, np.ndarray): - img_list = [img_list] - - # Calculate the aspect ratio of all text bars - width_list = [img.shape[1] / float(img.shape[0]) for img in img_list] - - # Sorting can speed up the recognition process - indices = np.argsort(np.array(width_list)) - - img_num = len(img_list) - rec_res = [['', 0.0]] * img_num - - batch_num = self.rec_batch_num - elapse = 0 - for beg_img_no in range(0, img_num, batch_num): - end_img_no = min(img_num, beg_img_no + batch_num) - max_wh_ratio = 0 - for ino in range(beg_img_no, end_img_no): - h, w = img_list[indices[ino]].shape[0:2] - wh_ratio = w * 1.0 / h - max_wh_ratio = max(max_wh_ratio, wh_ratio) - - norm_img_batch = [] - for ino in range(beg_img_no, end_img_no): - norm_img = self.resize_norm_img_svtr(img_list[indices[ino]]) - norm_img_batch.append(norm_img[np.newaxis, :]) - 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] - rec_result = self.postprocess_op(preds) - - for rno in range(len(rec_result)): - rec_res[indices[beg_img_no + rno]] = rec_result[rno] - elapse += time.time() - starttime - return rec_res, elapse - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument('--image_path', type=str, help='image_dir|image_path') - parser.add_argument('--config_path', type=str, default='config.yaml') - args = parser.parse_args() - - config = read_yaml(args.config_path) - text_recognizer = TextRecognizer(config) - - img = cv2.imread(args.image_path) - rec_res, predict_time = text_recognizer(img) - print(f'识别结果: {rec_res}\t cost: {predict_time}s') diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py deleted file mode 100644 index 9e0d436..0000000 --- a/python/rapidocr_onnxruntime/ch_ppocr_v3_rec/utils.py +++ /dev/null @@ -1,115 +0,0 @@ -# -*- 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])] - 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): - with open(yaml_path, 'rb') as f: - data = yaml.load(f, Loader=yaml.Loader) - return data - - -class CTCLabelDecode(object): - """ Convert between text-label and text-index """ - - def __init__(self, character_dict_path): - super(CTCLabelDecode, self).__init__() - - self.character_str = [] - assert character_dict_path is not None, "character_dict_path should not be None" - with open(character_dict_path, "rb") as fin: - lines = fin.readlines() - for line in lines: - line = line.decode('utf-8').strip("\n").strip("\r\n") - self.character_str.append(line) - self.character_str.append(' ') - - dict_character = self.add_special_char(self.character_str) - self.character = dict_character - - self.dict = {} - for i, char in enumerate(dict_character): - self.dict[char] = i - - def __call__(self, preds, label=None): - preds_idx = preds.argmax(axis=2) - preds_prob = preds.max(axis=2) - text = self.decode(preds_idx, preds_prob, - is_remove_duplicate=True) - if label is None: - return text - label = self.decode(label) - return text, label - - def add_special_char(self, dict_character): - dict_character = ['blank'] + dict_character - return dict_character - - def get_ignored_tokens(self): - return [0] # for ctc blank - - def decode(self, text_index, text_prob=None, is_remove_duplicate=False): - """ convert text-index into text-label. """ - - result_list = [] - ignored_tokens = self.get_ignored_tokens() - batch_size = len(text_index) - for batch_idx in range(batch_size): - char_list = [] - conf_list = [] - for idx in range(len(text_index[batch_idx])): - if text_index[batch_idx][idx] in ignored_tokens: - continue - if is_remove_duplicate: - # only for predict - if idx > 0 and text_index[batch_idx][idx - 1] == text_index[ - batch_idx][idx]: - continue - char_list.append(self.character[int(text_index[batch_idx][ - idx])]) - if text_prob is not None: - conf_list.append(text_prob[batch_idx][idx]) - else: - conf_list.append(1) - text = ''.join(char_list) - result_list.append((text, np.mean(conf_list))) - return result_list diff --git a/python/rapidocr_onnxruntime/rapid_ocr_api.py b/python/rapidocr_onnxruntime/rapid_ocr_api.py index f4b11a2..1bb7051 100644 --- a/python/rapidocr_onnxruntime/rapid_ocr_api.py +++ b/python/rapidocr_onnxruntime/rapid_ocr_api.py @@ -1,5 +1,6 @@ -# !/usr/bin/env python # -*- encoding: utf-8 -*- +# @Author: SWHL +# @Contact: liekkaskono@163.com import copy import importlib import sys @@ -14,8 +15,11 @@ sys.path.append(str(root_dir)) class TextSystem(object): - def __init__(self, config_path=str(root_dir / 'config.yaml')): + def __init__(self, config_path): super(TextSystem).__init__() + if not Path(config_path).exists(): + raise FileExistsError(f'{config_path} does not exist!') + config = self.read_yaml(config_path) global_config = config['Global'] diff --git a/python/tests/test_rec_v2.py b/python/tests/test_rec.py similarity index 100% rename from python/tests/test_rec_v2.py rename to python/tests/test_rec.py diff --git a/python/tests/test_rec_v3.py b/python/tests/test_rec_v3.py deleted file mode 100644 index 69fef63..0000000 --- a/python/tests/test_rec_v3.py +++ /dev/null @@ -1,24 +0,0 @@ -# -*- encoding: utf-8 -*- -# @Author: SWHL -# @Contact: liekkaskono@163.com -import base -import cv2 - - -module_name = 'ch_ppocr_v3_rec' -class_name = 'TextRecognizer' -TextRecognizer = base.init_module(module_name, class_name) - -yaml_path = base.root_dir / base.package_name / module_name / 'config.yaml' -config = base.read_yaml(str(yaml_path)) -config['model_path'] = str(base.root_dir / config['model_path']) -config['keys_path'] = str(base.root_dir / config['keys_path']) - -text_rec = TextRecognizer(config) - - -def test_rec(): - img_path = 'test_files/text_rec.jpg' - img = cv2.imread(img_path) - rec_res, elapse = text_rec([img]) - assert rec_res[0][0] == '韩国小馆'