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

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

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

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

L
LDOUBLEV 已提交
141 142 143
        return rec_res, predict_time


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


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