predict_rec.py 5.78 KB
Newer Older
LDOUBLEV's avatar
LDOUBLEV committed
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.
LDOUBLEV's avatar
LDOUBLEV committed
14
15
import os
import sys
WenmuZhou's avatar
WenmuZhou committed
16

17
__dir__ = os.path.dirname(os.path.abspath(__file__))
LDOUBLEV's avatar
LDOUBLEV committed
18
sys.path.append(__dir__)
19
sys.path.append(os.path.abspath(os.path.join(__dir__, '../..')))
LDOUBLEV's avatar
LDOUBLEV committed
20

21
22
os.environ["FLAGS_allocator_strategy"] = 'auto_growth'

LDOUBLEV's avatar
LDOUBLEV committed
23
24
25
26
import cv2
import numpy as np
import math
import time
WenmuZhou's avatar
WenmuZhou committed
27
import traceback
28
29

import tools.infer.utility as utility
WenmuZhou's avatar
WenmuZhou committed
30
31
from ppocr.postprocess import build_post_process
from ppocr.utils.logging import get_logger
32
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
LDOUBLEV's avatar
LDOUBLEV committed
33

WenmuZhou's avatar
WenmuZhou committed
34
35
logger = get_logger()

LDOUBLEV's avatar
LDOUBLEV committed
36
37
38

class TextRecognizer(object):
    def __init__(self, args):
39
        self.rec_image_shape = [int(v) for v in args.rec_image_shape.split(",")]
dyning's avatar
dyning committed
40
        self.character_type = args.rec_char_type
41
        self.rec_batch_num = args.rec_batch_num
tink2123's avatar
tink2123 committed
42
        self.rec_algorithm = args.rec_algorithm
WenmuZhou's avatar
WenmuZhou committed
43
44
        postprocess_params = {
            'name': 'CTCLabelDecode',
tink2123's avatar
tink2123 committed
45
            "character_type": args.rec_char_type,
46
            "character_dict_path": args.rec_char_dict_path,
WenmuZhou's avatar
WenmuZhou committed
47
            "use_space_char": args.use_space_char
tink2123's avatar
tink2123 committed
48
        }
WenmuZhou's avatar
WenmuZhou committed
49
50
51
        self.postprocess_op = build_post_process(postprocess_params)
        self.predictor, self.input_tensor, self.output_tensors = \
            utility.create_predictor(args, 'rec', logger)
LDOUBLEV's avatar
LDOUBLEV committed
52

53
    def resize_norm_img(self, img, max_wh_ratio):
LDOUBLEV's avatar
LDOUBLEV committed
54
        imgC, imgH, imgW = self.rec_image_shape
55
        assert imgC == img.shape[2]
56
        if self.character_type == "ch":
tink2123's avatar
tink2123 committed
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))
tink2123's avatar
tink2123 committed
64
        resized_image = cv2.resize(img, (resized_w, imgH))
LDOUBLEV's avatar
LDOUBLEV committed
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]))
zhangxin's avatar
zhangxin committed
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
WenmuZhou's avatar
WenmuZhou committed
85
        elapse = 0
LDOUBLEV's avatar
LDOUBLEV committed
86
87
88
        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
LDOUBLEV's avatar
LDOUBLEV committed
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)
tink2123's avatar
tink2123 committed
97
98
                norm_img = self.resize_norm_img(img_list[indices[ino]],
                                                max_wh_ratio)
LDOUBLEV's avatar
LDOUBLEV committed
99
100
101
102
103
                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()
WenmuZhou's avatar
WenmuZhou committed
104
105
            self.input_tensor.copy_from_cpu(norm_img_batch)
            self.predictor.run()
WenmuZhou's avatar
WenmuZhou committed
106
107
108
109
110
            outputs = []
            for output_tensor in self.output_tensors:
                output = output_tensor.copy_to_cpu()
                outputs.append(output)
            preds = outputs[0]
WenmuZhou's avatar
WenmuZhou committed
111
112
113
            rec_result = self.postprocess_op(preds)
            for rno in range(len(rec_result)):
                rec_res[indices[beg_img_no + rno]] = rec_result[rno]
114
            elapse += time.time() - starttime
WenmuZhou's avatar
WenmuZhou committed
115
        return rec_res, elapse
LDOUBLEV's avatar
LDOUBLEV committed
116
117


118
def main(args):
dyning's avatar
dyning committed
119
    image_file_list = get_image_file_list(args.image_dir)
LDOUBLEV's avatar
LDOUBLEV committed
120
121
122
123
    text_recognizer = TextRecognizer(args)
    valid_image_file_list = []
    img_list = []
    for image_file in image_file_list:
LDOUBLEV's avatar
LDOUBLEV committed
124
125
126
        img, flag = check_and_read_gif(image_file)
        if not flag:
            img = cv2.imread(image_file)
LDOUBLEV's avatar
LDOUBLEV committed
127
128
129
130
131
        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)
tink2123's avatar
tink2123 committed
132
133
    try:
        rec_res, predict_time = text_recognizer(img_list)
WenmuZhou's avatar
WenmuZhou committed
134
135
    except:
        logger.info(traceback.format_exc())
tink2123's avatar
tink2123 committed
136
        logger.info(
tink2123's avatar
tink2123 committed
137
138
139
140
            "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"
tink2123's avatar
tink2123 committed
141
            "Please set --rec_image_shape='3,32,100' and --rec_char_type='en' ")
tink2123's avatar
tink2123 committed
142
        exit()
LDOUBLEV's avatar
LDOUBLEV committed
143
    for ino in range(len(img_list)):
WenmuZhou's avatar
WenmuZhou committed
144
145
        logger.info("Predicts of {}:{}".format(valid_image_file_list[ino],
                                               rec_res[ino]))
WenmuZhou's avatar
WenmuZhou committed
146
    logger.info("Total predict time for {} images, cost: {:.3f}".format(
WenmuZhou's avatar
WenmuZhou committed
147
        len(img_list), predict_time))
148
149
150
151


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