add ch_ppocr_server_v2_rec

python/bpocr.py

@@ -220,19 +220,25 @@ if __name__ == '__main__':
     # v1.0
     # from ch_ppocr_mobile_v1_det import TextDetector
 
-    # v2.0
+    # v2.0 超轻量
+    # from ch_ppocr_mobile_v2_cls import TextClassifier
+    # from ch_ppocr_mobile_v2_det import TextDetector
+    # from ch_ppocr_mobile_v2_rec import TextRecognizer
+
+    # det_model_path = 'models/ch_ppocr_mobile_v2_det_train.onnx'
+    # cls_model_path = 'models/ch_ppocr_mobile_v2.0_cls_train.onnx'
+    # rec_model_path = 'models/ch_ppocr_mobile_v2.0_rec_pre.onnx'
+
+    # v2.0 通用模型
+    from ch_ppocr_server_v2_det import TextDetector
+    from ch_ppocr_server_v2_rec import TextRecognizer
     from ch_ppocr_mobile_v2_cls import TextClassifier
-    from ch_ppocr_mobile_v2_det import TextDetector
-    from ch_ppocr_mobile_v2_rec import TextRecognizer
 
-    # 通用模型路径
-    # from ch_ppocr_server_v2_det import TextDetector
-    # det_model_path = 'models/ch_ppocr_server_v2.0_det_train.onnx'
-
-    det_model_path = 'models/ch_ppocr_mobile_v2_det_train.onnx'
+    det_model_path = 'models/ch_ppocr_server_v2.0_det_train.onnx'
     cls_model_path = 'models/ch_ppocr_mobile_v2.0_cls_train.onnx'
-    rec_model_path = 'models/ch_ppocr_mobile_v2.0_rec_pre.onnx'
-    image_path = r'test_images/det_images/1.jpg'
+    rec_model_path = 'models/ch_ppocr_server_v2.0_rec_pre.onnx'
+
+    image_path = r'test_images/long1.jpg'
 
     text_sys = TextSystem(det_model_path, rec_model_path,
                           use_angle_cls=True,

python/ch_ppocr_server_v2_rec/__init__.py

@@ -0,0 +1,6 @@
+# !/usr/bin/env python
+# -*- encoding: utf-8 -*-
+# @File: __init__.py
+# @Time: 2021/02/03 08:15:19
+# @Author: Max
+from .text_recognize import TextRecognizer

python/ch_ppocr_server_v2_rec/text_recognize.py

@@ -0,0 +1,137 @@
+# 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.character_type = 'ch'
+        self.rec_batch_num = 6
+        self.rec_algorithm = 'CRNN'
+
+        self.character_dict_path = str(root_path / 'ppocr_keys_v1.txt')
+        self.use_space_char = True
+        self.postprocess_op = CTCLabelDecode(self.character_dict_path,
+                                             self.character_type,
+                                             self.use_space_char)
+        self.session = onnxruntime.InferenceSession(rec_model_path)
+
+    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]
+        if self.character_type == "ch":
+            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 = []
+        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[ino].shape[0:2]
+                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.copy()
+
+            starttime = time.time()
+            ort_inputs = {self.session.get_inputs()[0].name: norm_img_batch}
+            preds = self.session.run(None, ort_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_dir', 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_dir)
+    print(f'识别结果: {rec_res}\t cost: {predict_time}s')

python/ch_ppocr_server_v2_rec/utils.py

@@ -0,0 +1,124 @@
+# !/usr/bin/env python
+# -*- encoding: utf-8 -*-
+# @File: PostProcess.py
+# @Time: 2021/01/19 14:12:35
+# @Author: Max
+import cv2
+import os
+import imghdr
+import numpy as np
+
+
+class BaseRecLabelDecode(object):
+    """ Convert between text-label and text-index """
+
+    def __init__(self,
+                 character_dict_path=None,
+                 character_type='ch',
+                 use_space_char=False):
+        self.character_str = ""
+        assert character_dict_path is not None, "character_dict_path should not be None when character_type is ch"
+        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 += line
+        if use_space_char:
+            self.character_str += " "
+        dict_character = list(self.character_str)
+
+        self.character_type = character_type
+        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 add_special_char(self, dict_character):
+        return dict_character
+
+    def decode(self, text_index, text_prob=None, is_remove_duplicate=True):
+        """ 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
+
+    def get_ignored_tokens(self):
+        return [0]  # for ctc blank
+
+
+class CTCLabelDecode(BaseRecLabelDecode):
+    """ Convert between text-label and text-index """
+
+    def __init__(self,
+                 character_dict_path=None,
+                 character_type='ch',
+                 use_space_char=False,
+                 **kwargs):
+        super(CTCLabelDecode, self).__init__(character_dict_path,
+                                             character_type, use_space_char)
+
+    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)
+        if label is None:
+            return text
+        label = self.decode(label, is_remove_duplicate=False)
+        return text, label
+
+    def add_special_char(self, dict_character):
+        dict_character = ['blank'] + dict_character
+        return dict_character
+
+
+def get_image_file_list(img_file):
+    imgs_lists = []
+    if img_file is None or not os.path.exists(img_file):
+        raise Exception("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 Exception("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

python/readme.md

@@ -2,8 +2,8 @@
 
 - [x] 超轻量2.0方向分类器模型整理
 - [x] 通用2.0文本检测模型整理
-- [ ] 通用2.0文本识别模型整理
-- [ ] 通用2.0方向分类器模型整理
+- [x] 通用2.0文本识别模型整理
+- [x] ~~通用2.0方向分类器模型整理~~ → 同**超轻量2.0方向分类器模型整理**
 
 ### 检测、方向分类和识别模型下载：