diff --git a/python/README.md b/python/README.md index f7be42b..6eb69c6 100755 --- a/python/README.md +++ b/python/README.md @@ -1,19 +1,3 @@ -### 🎉2021-06-10 update: -- 添加server版文本识别模型，详情见[提取码：30jv](https://pan.baidu.com/s/1qkqWK4wRdMjqGGbzR-FyWg) - -### 2021-06-09 update: -- 添加字体文件到模型下载链接中 - -### 2021-05-22 update: -- 添加日文识别推理代码和相应模型，详情见`japan_ppocr_mobile_v2_rec`目录 -- 整理代码，去除无效代码 -- 每个单独模块，添加执行脚本示例 - -### 2021-03-24 update: -- 更改文本识别模型支持ort1.7,推理速度有着较大提升 - ---- - ### 运行 1. 下载相应模型和用于显示的字体文件 - [提取码：30jv](https://pan.baidu.com/s/1qkqWK4wRdMjqGGbzR-FyWg) diff --git a/python/en_number_ppocr_mobile_v2_rec/__init__.py b/python/en_number_ppocr_mobile_v2_rec/__init__.py new file mode 100644 index 0000000..686ffdd --- /dev/null +++ b/python/en_number_ppocr_mobile_v2_rec/__init__.py @@ -0,0 +1 @@ +from .text_recognize import TextRecognizer \ No newline at end of file diff --git a/python/en_number_ppocr_mobile_v2_rec/en_dict.txt b/python/en_number_ppocr_mobile_v2_rec/en_dict.txt new file mode 100644 index 0000000..7677d31 --- /dev/null +++ b/python/en_number_ppocr_mobile_v2_rec/en_dict.txt @@ -0,0 +1,95 @@ +0 +1 +2 +3 +4 +5 +6 +7 +8 +9 +: +; +< += +> +? +@ +A +B +C +D +E +F +G +H +I +J +K +L +M +N +O +P +Q +R +S +T +U +V +W +X +Y +Z +[ +\ +] +^ +_ +` +a +b +c +d +e +f +g +h +i +j +k +l +m +n +o +p +q +r +s +t +u +v +w +x +y +z +{ +| +} +~ +! +" +# +$ +% +& +' +( +) +* ++ +, +- +. +/ + diff --git a/python/en_number_ppocr_mobile_v2_rec/run.sh b/python/en_number_ppocr_mobile_v2_rec/run.sh new file mode 100644 index 0000000..f519b8b --- /dev/null +++ b/python/en_number_ppocr_mobile_v2_rec/run.sh @@ -0,0 +1,9 @@ +#! /bin/bash +set -e errexit + +root_path=$(cd $(dirname $0); pwd) +image_path=${root_path}"/test_images/rec_images/2021-01-18_14-22-32.png" +model_path=${root_path}"/models/ch_ppocr_mobile_v2.0_rec_infer.onnx" + +python text_recognize.py --image_path ${image_path} \ + --model_path ${model_path} \ No newline at end of file diff --git a/python/en_number_ppocr_mobile_v2_rec/text_recognize.py b/python/en_number_ppocr_mobile_v2_rec/text_recognize.py new file mode 100644 index 0000000..a6a1357 --- /dev/null +++ b/python/en_number_ppocr_mobile_v2_rec/text_recognize.py @@ -0,0 +1,136 @@ +# 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 + +import cv2 +import numpy as np +import onnxruntime + +try: + from .utils import CTCLabelDecode, check_and_read_gif, get_image_file_list +except: + from utils import CTCLabelDecode, check_and_read_gif, get_image_file_list + + +root_path = Path(__file__).parent + + +class TextRecognizer(object): + def __init__(self, rec_model_path): + self.rec_image_shape = [3, 32, 320] + self.rec_batch_num = 6 + + self.character_dict_path = str(root_path / 'en_dict.txt') + 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(rec_model_path, sess_opt) + + 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 + + @staticmethod + def load_image(image_dir): + image_file_list = get_image_file_list(image_dir) + valid_image_file_list = [] + img_list = [] + for image_file in image_file_list: + img, flag = check_and_read_gif(image_file) + if not flag: + img = cv2.imread(image_file) + if img is None: + print(f"error in loading image:{image_file}") + continue + valid_image_file_list.append(image_file) + img_list.append(img) + return img_list + + def __call__(self, image_dir): + if not isinstance(image_dir, list): + img_list = self.load_image(image_dir) + else: + img_list = image_dir + img_num = len(img_list) + # Calculate the aspect ratio of all text bars + width_list = [] + for img in img_list: + width_list.append(img.shape[1] / float(img.shape[0])) + # Sorting can speed up the recognition process + indices = np.argsort(np.array(width_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) + norm_img_batch = [] + 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) + for ino in range(beg_img_no, end_img_no): + norm_img = self.resize_norm_img(img_list[indices[ino]], + max_wh_ratio) + norm_img = norm_img[np.newaxis, :] + norm_img_batch.append(norm_img) + norm_img_batch = np.concatenate(norm_img_batch) + norm_img_batch = norm_img_batch.astype(np.float32) + + # onnx模型推理 + starttime = time.time() + onnx_inputs = {self.session.get_inputs()[0].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('--model_path', type=str, help='rec_model_path') + args = parser.parse_args() + + text_recognizer = TextRecognizer(args.model_path) + rec_res, predict_time = text_recognizer(args.image_path) + print(f'识别结果: {rec_res}\t cost: {predict_time}s') diff --git a/python/en_number_ppocr_mobile_v2_rec/utils.py b/python/en_number_ppocr_mobile_v2_rec/utils.py new file mode 100644 index 0000000..cb8c8d9 --- /dev/null +++ b/python/en_number_ppocr_mobile_v2_rec/utils.py @@ -0,0 +1,109 @@ +# !/usr/bin/env python +# -*- encoding: utf-8 -*- +# @File: utils.py +# @Time: 2021/05/22 10:53:17 +# @Author: Max +import cv2 +import os +import imghdr +import numpy as np + + +def get_image_file_list(img_file): + imgs_lists = [] + if img_file is None or not os.path.exists(img_file): + raise FileNotFoundError("not found any img file in {}".format(img_file)) + + img_end = {'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif', 'GIF'} + if os.path.isfile(img_file) and imghdr.what(img_file) in img_end: + imgs_lists.append(img_file) + elif os.path.isdir(img_file): + for single_file in os.listdir(img_file): + file_path = os.path.join(img_file, single_file) + if imghdr.what(file_path) in img_end: + imgs_lists.append(file_path) + if len(imgs_lists) == 0: + raise FileNotFoundError("not found any img file in {}".format(img_file)) + return imgs_lists + + +def check_and_read_gif(img_path): + if os.path.basename(img_path)[-3:] in ['gif', 'GIF']: + gif = cv2.VideoCapture(img_path) + ret, frame = gif.read() + if not ret: + print("Cannot read {}. This gif image maybe corrupted.") + return None, False + if len(frame.shape) == 2 or frame.shape[-1] == 1: + frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB) + imgvalue = frame[:, :, ::-1] + return imgvalue, True + return None, False + + +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.character_str + + dict_character = self.add_special_char(dict_character) + self.dict = {} + for i, char in enumerate(dict_character): + self.dict[char] = i + self.character = dict_character + + def __call__(self, preds, label=None, *args, **kwargs): + 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/test_images/det_images/2021-06-06_22-24-36.png b/python/test_images/det_images/2021-06-06_22-24-36.png new file mode 100644 index 0000000..d613cc4 Binary files /dev/null and b/python/test_images/det_images/2021-06-06_22-24-36.png differ diff --git a/python/test_images/rec_images/2021-01-18_14-22-32.png b/python/test_images/rec_images/imgs_words_ch/1.png old mode 100755 new mode 100644 similarity index 100% rename from python/test_images/rec_images/2021-01-18_14-22-32.png rename to python/test_images/rec_images/imgs_words_ch/1.png diff --git a/python/test_images/rec_images/2021-01-19_13-44-34.png b/python/test_images/rec_images/imgs_words_ch/2.png old mode 100755 new mode 100644 similarity index 100% rename from python/test_images/rec_images/2021-01-19_13-44-34.png rename to python/test_images/rec_images/imgs_words_ch/2.png diff --git a/python/test_images/rec_images/imgs_words_en/word_10.png b/python/test_images/rec_images/imgs_words_en/word_10.png new file mode 100644 index 0000000..07370f7 Binary files /dev/null and b/python/test_images/rec_images/imgs_words_en/word_10.png differ diff --git a/python/test_images/rec_images/japan/1.jpg b/python/test_images/rec_images/imgs_words_japan/1.jpg similarity index 100% rename from python/test_images/rec_images/japan/1.jpg rename to python/test_images/rec_images/imgs_words_japan/1.jpg