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
        char_ops_params = {
            "character_type": args.rec_char_type,
42
            "character_dict_path": args.rec_char_dict_path,
T
tink2123 已提交
43
            "use_space_char": args.use_space_char
T
tink2123 已提交
44
        }
T
tink2123 已提交
45 46
        if self.rec_algorithm != "RARE":
            char_ops_params['loss_type'] = 'ctc'
T
tink2123 已提交
47
            self.loss_type = 'ctc'
T
tink2123 已提交
48 49
        else:
            char_ops_params['loss_type'] = 'attention'
T
tink2123 已提交
50
            self.loss_type = 'attention'
L
LDOUBLEV 已提交
51 52
        self.char_ops = CharacterOps(char_ops_params)

53
    def resize_norm_img(self, img, max_wh_ratio):
L
LDOUBLEV 已提交
54
        imgC, imgH, imgW = self.rec_image_shape
55
        assert imgC == img.shape[2]
56
        if self.character_type == "ch":
T
tink2123 已提交
57
            imgW = int((32 * max_wh_ratio))
58
        h, w = img.shape[:2]
59 60 61 62 63
        ratio = w / float(h)
        if math.ceil(imgH * ratio) > imgW:
            resized_w = imgW
        else:
            resized_w = int(math.ceil(imgH * ratio))
T
tink2123 已提交
64
        resized_image = cv2.resize(img, (resized_w, imgH))
L
LDOUBLEV 已提交
65 66 67 68 69 70 71 72 73 74
        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)
75
        # Calculate the aspect ratio of all text bars
76 77 78
        width_list = []
        for img in img_list:
            width_list.append(img.shape[1] / float(img.shape[0]))
张欣-男's avatar
张欣-男 已提交
79
        # Sorting can speed up the recognition process
80 81 82 83
        indices = np.argsort(np.array(width_list))

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

T
tink2123 已提交
107
            if self.loss_type == "ctc":
T
tink2123 已提交
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125
                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]])
L
fix bug  
LDOUBLEV 已提交
126
                    if len(valid_ind) == 0:
127
                        continue
128 129
                    # rec_res.append([preds_text, score])
                    rec_res[indices[beg_img_no + rno]] = [preds_text, score]
T
tink2123 已提交
130 131 132
            else:
                rec_idx_batch = self.output_tensors[0].copy_to_cpu()
                predict_batch = self.output_tensors[1].copy_to_cpu()
T
tink2123 已提交
133 134
                elapse = time.time() - starttime
                predict_time += elapse
T
tink2123 已提交
135 136 137 138 139 140 141 142 143
                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)
144 145
                    # rec_res.append([preds_text, score])
                    rec_res[indices[beg_img_no + rno]] = [preds_text, score]
T
tink2123 已提交
146

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


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


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