predict_rec.py 7.3 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 56
        if self.character_type == "ch":
            imgW = int(math.ceil(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
                    # rec_res.append([preds_text, score])
                    rec_res[indices[beg_img_no + rno]] = [preds_text, score]
T
tink2123 已提交
127 128 129
            else:
                rec_idx_batch = self.output_tensors[0].copy_to_cpu()
                predict_batch = self.output_tensors[1].copy_to_cpu()
T
tink2123 已提交
130 131
                elapse = time.time() - starttime
                predict_time += elapse
T
tink2123 已提交
132 133 134 135 136 137 138 139 140
                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)
141 142
                    # rec_res.append([preds_text, score])
                    rec_res[indices[beg_img_no + rno]] = [preds_text, score]
T
tink2123 已提交
143

L
LDOUBLEV 已提交
144 145 146
        return rec_res, predict_time


147
def main(args):
D
dyning 已提交
148
    image_file_list = get_image_file_list(args.image_dir)
L
LDOUBLEV 已提交
149 150 151 152
    text_recognizer = TextRecognizer(args)
    valid_image_file_list = []
    img_list = []
    for image_file in image_file_list:
153
        img = cv2.imread(image_file, cv2.IMREAD_COLOR)
L
LDOUBLEV 已提交
154 155 156 157 158
        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 已提交
159 160
    try:
        rec_res, predict_time = text_recognizer(img_list)
T
tink2123 已提交
161 162
    except Exception as e:
        print(e)
T
tink2123 已提交
163
        logger.info(
T
tink2123 已提交
164 165 166 167
            "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 已提交
168
            "Please set --rec_image_shape='3,32,100' and --rec_char_type='en' ")
T
tink2123 已提交
169
        exit()
L
LDOUBLEV 已提交
170 171 172 173
    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))
174 175 176 177


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