predict_rec.py 7.4 KB
Newer Older
L
LDOUBLEV 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13
# 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.
L
LDOUBLEV 已提交
14 15
import os
import sys
16
__dir__ = os.path.dirname(os.path.abspath(__file__))
L
LDOUBLEV 已提交
17
sys.path.append(__dir__)
18
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
L
LDOUBLEV 已提交
19

L
LDOUBLEV 已提交
20
import tools.infer.utility as utility
L
LDOUBLEV 已提交
21 22
from ppocr.utils.utility import initial_logger
logger = initial_logger()
D
dyning 已提交
23
from ppocr.utils.utility import get_image_file_list
L
LDOUBLEV 已提交
24 25 26 27 28 29 30 31 32 33 34 35
import cv2
import copy
import numpy as np
import math
import time
from ppocr.utils.character import CharacterOps


class TextRecognizer(object):
    def __init__(self, args):
        self.predictor, self.input_tensor, self.output_tensors =\
            utility.create_predictor(args, mode="rec")
36
        self.rec_image_shape = [int(v) for v in args.rec_image_shape.split(",")]
D
dyning 已提交
37
        self.character_type = args.rec_char_type
38
        self.rec_batch_num = args.rec_batch_num
T
tink2123 已提交
39
        self.rec_algorithm = args.rec_algorithm
T
tink2123 已提交
40 41 42 43
        char_ops_params = {
            "character_type": args.rec_char_type,
            "character_dict_path": args.rec_char_dict_path
        }
T
tink2123 已提交
44 45
        if self.rec_algorithm != "RARE":
            char_ops_params['loss_type'] = 'ctc'
T
tink2123 已提交
46
            self.loss_type = 'ctc'
T
tink2123 已提交
47 48
        else:
            char_ops_params['loss_type'] = 'attention'
T
tink2123 已提交
49
            self.loss_type = 'attention'
L
LDOUBLEV 已提交
50 51
        self.char_ops = CharacterOps(char_ops_params)

52
    def resize_norm_img(self, img, max_wh_ratio):
L
LDOUBLEV 已提交
53
        imgC, imgH, imgW = self.rec_image_shape
54
        assert imgC == img.shape[2]
55
        if self.character_type == "ch":
T
tink2123 已提交
56
            imgW = int((32 * max_wh_ratio))
57
        h, w = img.shape[:2]
58 59 60 61 62
        ratio = w / float(h)
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))
T
tink2123 已提交
63
        resized_image = cv2.resize(img, (resized_w, imgH))
L
LDOUBLEV 已提交
64 65 66 67 68 69 70 71 72 73
        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((imgC, imgH, imgW), dtype=np.float32)
        padding_im[:, :, 0:resized_w] = resized_image
        return padding_im

    def __call__(self, img_list):
        img_num = len(img_list)
74
        # Calculate the aspect ratio of all text bars
75 76 77
        width_list = []
        for img in img_list:
            width_list.append(img.shape[1] / float(img.shape[0]))
张欣-男's avatar
张欣-男 已提交
78
        # Sorting can speed up the recognition process
79 80 81 82
        indices = np.argsort(np.array(width_list))

        # rec_res = []
        rec_res = [['', 0.0]] * img_num
83
        batch_num = self.rec_batch_num
L
LDOUBLEV 已提交
84 85 86 87
        predict_time = 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 = []
88
            max_wh_ratio = 0
L
LDOUBLEV 已提交
89
            for ino in range(beg_img_no, end_img_no):
90 91
                # h, w = img_list[ino].shape[0:2]
                h, w = img_list[indices[ino]].shape[0:2]
92 93 94
                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):
95
                # norm_img = self.resize_norm_img(img_list[ino], max_wh_ratio)
T
tink2123 已提交
96 97
                norm_img = self.resize_norm_img(img_list[indices[ino]],
                                                max_wh_ratio)
L
LDOUBLEV 已提交
98 99 100 101 102 103 104
                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()
            self.input_tensor.copy_from_cpu(norm_img_batch)
            self.predictor.zero_copy_run()
T
tink2123 已提交
105

T
tink2123 已提交
106
            if self.loss_type == "ctc":
T
tink2123 已提交
107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124
                rec_idx_batch = self.output_tensors[0].copy_to_cpu()
                rec_idx_lod = self.output_tensors[0].lod()[0]
                predict_batch = self.output_tensors[1].copy_to_cpu()
                predict_lod = self.output_tensors[1].lod()[0]
                elapse = time.time() - starttime
                predict_time += elapse
                for rno in range(len(rec_idx_lod) - 1):
                    beg = rec_idx_lod[rno]
                    end = rec_idx_lod[rno + 1]
                    rec_idx_tmp = rec_idx_batch[beg:end, 0]
                    preds_text = self.char_ops.decode(rec_idx_tmp)
                    beg = predict_lod[rno]
                    end = predict_lod[rno + 1]
                    probs = predict_batch[beg:end, :]
                    ind = np.argmax(probs, axis=1)
                    blank = probs.shape[1]
                    valid_ind = np.where(ind != (blank - 1))[0]
                    score = np.mean(probs[valid_ind, ind[valid_ind]])
125 126
                    if not valid_ind:
                        continue
127 128
                    # rec_res.append([preds_text, score])
                    rec_res[indices[beg_img_no + rno]] = [preds_text, score]
T
tink2123 已提交
129 130 131
            else:
                rec_idx_batch = self.output_tensors[0].copy_to_cpu()
                predict_batch = self.output_tensors[1].copy_to_cpu()
T
tink2123 已提交
132 133
                elapse = time.time() - starttime
                predict_time += elapse
T
tink2123 已提交
134 135 136 137 138 139 140 141 142
                for rno in range(len(rec_idx_batch)):
                    end_pos = np.where(rec_idx_batch[rno, :] == 1)[0]
                    if len(end_pos) <= 1:
                        preds = rec_idx_batch[rno, 1:]
                        score = np.mean(predict_batch[rno, 1:])
                    else:
                        preds = rec_idx_batch[rno, 1:end_pos[1]]
                        score = np.mean(predict_batch[rno, 1:end_pos[1]])
                    preds_text = self.char_ops.decode(preds)
143 144
                    # rec_res.append([preds_text, score])
                    rec_res[indices[beg_img_no + rno]] = [preds_text, score]
T
tink2123 已提交
145

L
LDOUBLEV 已提交
146 147 148
        return rec_res, predict_time


149
def main(args):
D
dyning 已提交
150
    image_file_list = get_image_file_list(args.image_dir)
L
LDOUBLEV 已提交
151 152 153 154
    text_recognizer = TextRecognizer(args)
    valid_image_file_list = []
    img_list = []
    for image_file in image_file_list:
155
        img = cv2.imread(image_file, cv2.IMREAD_COLOR)
L
LDOUBLEV 已提交
156 157 158 159 160
        if img is None:
            logger.info("error in loading image:{}".format(image_file))
            continue
        valid_image_file_list.append(image_file)
        img_list.append(img)
T
tink2123 已提交
161 162
    try:
        rec_res, predict_time = text_recognizer(img_list)
T
tink2123 已提交
163 164
    except Exception as e:
        print(e)
T
tink2123 已提交
165
        logger.info(
T
tink2123 已提交
166 167 168 169
            "ERROR!!!! \n"
            "Please read the FAQ:https://github.com/PaddlePaddle/PaddleOCR#faq \n"
            "If your model has tps module:  "
            "TPS does not support variable shape.\n"
T
tink2123 已提交
170
            "Please set --rec_image_shape='3,32,100' and --rec_char_type='en' ")
T
tink2123 已提交
171
        exit()
L
LDOUBLEV 已提交
172 173 174 175
    for ino in range(len(img_list)):
        print("Predicts of %s:%s" % (valid_image_file_list[ino], rec_res[ino]))
    print("Total predict time for %d images:%.3f" %
          (len(img_list), predict_time))
176 177 178 179


if __name__ == "__main__":
    main(utility.parse_args())