predict_system.py 10 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.
14
15
import os
import sys
LDOUBLEV's avatar
LDOUBLEV committed
16
import subprocess
WenmuZhou's avatar
WenmuZhou committed
17

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

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

LDOUBLEV's avatar
LDOUBLEV committed
24
25
26
27
import cv2
import copy
import numpy as np
import time
WenmuZhou's avatar
WenmuZhou committed
28
import logging
LDOUBLEV's avatar
LDOUBLEV committed
29
from PIL import Image
WenmuZhou's avatar
WenmuZhou committed
30
31
32
import tools.infer.utility as utility
import tools.infer.predict_rec as predict_rec
import tools.infer.predict_det as predict_det
WenmuZhou's avatar
WenmuZhou committed
33
import tools.infer.predict_cls as predict_cls
WenmuZhou's avatar
WenmuZhou committed
34
35
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
from ppocr.utils.logging import get_logger
LDOUBLEV's avatar
LDOUBLEV committed
36
37
from tools.infer.utility import draw_ocr_box_txt, get_current_memory_mb
import tools.infer.benchmark_utils as benchmark_utils
WenmuZhou's avatar
WenmuZhou committed
38
39
logger = get_logger()

LDOUBLEV's avatar
LDOUBLEV committed
40
41
42

class TextSystem(object):
    def __init__(self, args):
WenmuZhou's avatar
WenmuZhou committed
43
44
45
        if not args.show_log:
            logger.setLevel(logging.INFO)

LDOUBLEV's avatar
LDOUBLEV committed
46
47
        self.text_detector = predict_det.TextDetector(args)
        self.text_recognizer = predict_rec.TextRecognizer(args)
WenmuZhou's avatar
WenmuZhou committed
48
        self.use_angle_cls = args.use_angle_cls
WenmuZhou's avatar
WenmuZhou committed
49
        self.drop_score = args.drop_score
WenmuZhou's avatar
WenmuZhou committed
50
51
        if self.use_angle_cls:
            self.text_classifier = predict_cls.TextClassifier(args)
LDOUBLEV's avatar
LDOUBLEV committed
52
53

    def get_rotate_crop_image(self, img, points):
54
        '''
LDOUBLEV's avatar
LDOUBLEV committed
55
56
57
58
59
60
61
62
        img_height, img_width = img.shape[0:2]
        left = int(np.min(points[:, 0]))
        right = int(np.max(points[:, 0]))
        top = int(np.min(points[:, 1]))
        bottom = int(np.max(points[:, 1]))
        img_crop = img[top:bottom, left:right, :].copy()
        points[:, 0] = points[:, 0] - left
        points[:, 1] = points[:, 1] - top
63
        '''
LDOUBLEV's avatar
LDOUBLEV committed
64
65
66
67
68
69
70
71
72
        img_crop_width = int(
            max(
                np.linalg.norm(points[0] - points[1]),
                np.linalg.norm(points[2] - points[3])))
        img_crop_height = int(
            max(
                np.linalg.norm(points[0] - points[3]),
                np.linalg.norm(points[1] - points[2])))
        pts_std = np.float32([[0, 0], [img_crop_width, 0],
73
74
                              [img_crop_width, img_crop_height],
                              [0, img_crop_height]])
LDOUBLEV's avatar
LDOUBLEV committed
75
        M = cv2.getPerspectiveTransform(points, pts_std)
LDOUBLEV's avatar
LDOUBLEV committed
76
77
78
79
80
        dst_img = cv2.warpPerspective(
            img,
            M, (img_crop_width, img_crop_height),
            borderMode=cv2.BORDER_REPLICATE,
            flags=cv2.INTER_CUBIC)
LDOUBLEV's avatar
LDOUBLEV committed
81
82
83
84
85
86
87
88
89
        dst_img_height, dst_img_width = dst_img.shape[0:2]
        if dst_img_height * 1.0 / dst_img_width >= 1.5:
            dst_img = np.rot90(dst_img)
        return dst_img

    def print_draw_crop_rec_res(self, img_crop_list, rec_res):
        bbox_num = len(img_crop_list)
        for bno in range(bbox_num):
            cv2.imwrite("./output/img_crop_%d.jpg" % bno, img_crop_list[bno])
WenmuZhou's avatar
WenmuZhou committed
90
            logger.info(bno, rec_res[bno])
LDOUBLEV's avatar
LDOUBLEV committed
91

92
    def __call__(self, img, cls=True):
LDOUBLEV's avatar
LDOUBLEV committed
93
94
        ori_im = img.copy()
        dt_boxes, elapse = self.text_detector(img)
LDOUBLEV's avatar
LDOUBLEV committed
95

WenmuZhou's avatar
WenmuZhou committed
96
        logger.debug("dt_boxes num : {}, elapse : {}".format(
WenmuZhou's avatar
WenmuZhou committed
97
            len(dt_boxes), elapse))
LDOUBLEV's avatar
LDOUBLEV committed
98
99
100
        if dt_boxes is None:
            return None, None
        img_crop_list = []
101
102
103

        dt_boxes = sorted_boxes(dt_boxes)

LDOUBLEV's avatar
LDOUBLEV committed
104
105
106
107
        for bno in range(len(dt_boxes)):
            tmp_box = copy.deepcopy(dt_boxes[bno])
            img_crop = self.get_rotate_crop_image(ori_im, tmp_box)
            img_crop_list.append(img_crop)
108
        if self.use_angle_cls and cls:
WenmuZhou's avatar
WenmuZhou committed
109
110
            img_crop_list, angle_list, elapse = self.text_classifier(
                img_crop_list)
WenmuZhou's avatar
WenmuZhou committed
111
            logger.debug("cls num  : {}, elapse : {}".format(
WenmuZhou's avatar
WenmuZhou committed
112
113
                len(img_crop_list), elapse))

LDOUBLEV's avatar
LDOUBLEV committed
114
        rec_res, elapse = self.text_recognizer(img_crop_list)
WenmuZhou's avatar
WenmuZhou committed
115
        logger.debug("rec_res num  : {}, elapse : {}".format(
WenmuZhou's avatar
WenmuZhou committed
116
            len(rec_res), elapse))
117
        # self.print_draw_crop_rec_res(img_crop_list, rec_res)
WenmuZhou's avatar
WenmuZhou committed
118
119
120
121
122
123
124
        filter_boxes, filter_rec_res = [], []
        for box, rec_reuslt in zip(dt_boxes, rec_res):
            text, score = rec_reuslt
            if score >= self.drop_score:
                filter_boxes.append(box)
                filter_rec_res.append(rec_reuslt)
        return filter_boxes, filter_rec_res
LDOUBLEV's avatar
LDOUBLEV committed
125
126


127
128
129
130
def sorted_boxes(dt_boxes):
    """
    Sort text boxes in order from top to bottom, left to right
    args:
tink2123's avatar
tink2123 committed
131
        dt_boxes(array):detected text boxes with shape [4, 2]
132
133
134
135
    return:
        sorted boxes(array) with shape [4, 2]
    """
    num_boxes = dt_boxes.shape[0]
136
    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
137
138
139
    _boxes = list(sorted_boxes)

    for i in range(num_boxes - 1):
WenmuZhou's avatar
WenmuZhou committed
140
141
        if abs(_boxes[i + 1][0][1] - _boxes[i][0][1]) < 10 and \
                (_boxes[i + 1][0][0] < _boxes[i][0][0]):
142
143
144
145
146
147
            tmp = _boxes[i]
            _boxes[i] = _boxes[i + 1]
            _boxes[i + 1] = tmp
    return _boxes


148
def main(args):
LDOUBLEV's avatar
LDOUBLEV committed
149
    image_file_list = get_image_file_list(args.image_dir)
LDOUBLEV's avatar
LDOUBLEV committed
150
    image_file_list = image_file_list[args.process_id::args.total_process_num]
LDOUBLEV's avatar
LDOUBLEV committed
151
    text_sys = TextSystem(args)
LDOUBLEV's avatar
LDOUBLEV committed
152
    is_visualize = True
WenmuZhou's avatar
WenmuZhou committed
153
    font_path = args.vis_font_path
WenmuZhou's avatar
WenmuZhou committed
154
    drop_score = args.drop_score
LDOUBLEV's avatar
LDOUBLEV committed
155
156
157
158
159
    total_time = 0
    cpu_mem, gpu_mem, gpu_util = 0, 0, 0
    _st = time.time()
    count = 0
    for idx, image_file in enumerate(image_file_list):
LDOUBLEV's avatar
LDOUBLEV committed
160
161
162
        img, flag = check_and_read_gif(image_file)
        if not flag:
            img = cv2.imread(image_file)
LDOUBLEV's avatar
LDOUBLEV committed
163
        if img is None:
164
            logger.info("error in loading image:{}".format(image_file))
LDOUBLEV's avatar
LDOUBLEV committed
165
166
167
168
            continue
        starttime = time.time()
        dt_boxes, rec_res = text_sys(img)
        elapse = time.time() - starttime
LDOUBLEV's avatar
LDOUBLEV committed
169
170
171
172
173
174
175
        total_time += elapse
        if args.benchmark and idx % 20 == 0:
            cm, gm, gu = get_current_memory_mb(0)
            cpu_mem += cm
            gpu_mem += gm
            gpu_util += gu
            count += 1
LDOUBLEV's avatar
LDOUBLEV committed
176

LDOUBLEV's avatar
LDOUBLEV committed
177
178
        logger.info(
            str(idx) + "  Predict time of %s: %.3fs" % (image_file, elapse))
WenmuZhou's avatar
WenmuZhou committed
179
180
        for text, score in rec_res:
            logger.info("{}, {:.3f}".format(text, score))
LDOUBLEV's avatar
LDOUBLEV committed
181
182
183
184
185
186
187

        if is_visualize:
            image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
            boxes = dt_boxes
            txts = [rec_res[i][0] for i in range(len(rec_res))]
            scores = [rec_res[i][1] for i in range(len(rec_res))]

WenmuZhou's avatar
WenmuZhou committed
188
189
190
191
192
193
194
            draw_img = draw_ocr_box_txt(
                image,
                boxes,
                txts,
                scores,
                drop_score=drop_score,
                font_path=font_path)
195
            draw_img_save = "./inference_results/"
LDOUBLEV's avatar
LDOUBLEV committed
196
197
            if not os.path.exists(draw_img_save):
                os.makedirs(draw_img_save)
LDOUBLEV's avatar
LDOUBLEV committed
198
199
            if flag:
                image_file = image_file[:-3] + "png"
LDOUBLEV's avatar
LDOUBLEV committed
200
201
            cv2.imwrite(
                os.path.join(draw_img_save, os.path.basename(image_file)),
dyning's avatar
dyning committed
202
                draw_img[:, :, ::-1])
WenmuZhou's avatar
WenmuZhou committed
203
            logger.info("The visualized image saved in {}".format(
204
                os.path.join(draw_img_save, os.path.basename(image_file))))
205

LDOUBLEV's avatar
LDOUBLEV committed
206
207
    logger.info("The predict total time is {}".format(time.time() - _st))
    logger.info("\nThe predict total time is {}".format(total_time))
208

LDOUBLEV's avatar
LDOUBLEV committed
209
210
211
212
213
214
215
    img_num = text_sys.text_detector.det_times.img_num
    if args.benchmark:
        mems = {
            'cpu_rss_mb': cpu_mem / count,
            'gpu_rss_mb': gpu_mem / count,
            'gpu_util': gpu_util * 100 / count
        }
littletomatodonkey's avatar
littletomatodonkey committed
216
    else:
LDOUBLEV's avatar
LDOUBLEV committed
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
        mems = None
    det_time_dict = text_sys.text_detector.det_times.report(average=True)
    rec_time_dict = text_sys.text_recognizer.rec_times.report(average=True)
    det_model_name = args.det_model_dir
    rec_model_name = args.rec_model_dir

    # construct det log information
    model_info = {
        'model_name': args.det_model_dir.split('/')[-1],
        'precision': args.precision
    }
    data_info = {
        'batch_size': 1,
        'shape': 'dynamic_shape',
        'data_num': det_time_dict['img_num']
    }
    perf_info = {
        'preprocess_time_s': det_time_dict['preprocess_time'],
        'inference_time_s': det_time_dict['inference_time'],
        'postprocess_time_s': det_time_dict['postprocess_time'],
        'total_time_s': det_time_dict['total_time']
    }

    benchmark_log = benchmark_utils.PaddleInferBenchmark(
        text_sys.text_detector.config, model_info, data_info, perf_info, mems,
        args.save_log_path)
    benchmark_log("Det")

    # construct rec log information
    model_info = {
        'model_name': args.rec_model_dir.split('/')[-1],
        'precision': args.precision
    }
    data_info = {
        'batch_size': args.rec_batch_num,
        'shape': 'dynamic_shape',
        'data_num': rec_time_dict['img_num']
    }
    perf_info = {
        'preprocess_time_s': rec_time_dict['preprocess_time'],
        'inference_time_s': rec_time_dict['inference_time'],
        'postprocess_time_s': rec_time_dict['postprocess_time'],
        'total_time_s': rec_time_dict['total_time']
    }
    benchmark_log = benchmark_utils.PaddleInferBenchmark(
        text_sys.text_recognizer.config, model_info, data_info, perf_info, mems,
        args.save_log_path)
    benchmark_log("Rec")


if __name__ == "__main__":
LDOUBLEV's avatar
LDOUBLEV committed
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
    args = utility.parse_args()
    if args.use_mp:
        p_list = []
        total_process_num = args.total_process_num
        for process_id in range(total_process_num):
            cmd = [sys.executable, "-u"] + sys.argv + [
                "--process_id={}".format(process_id),
                "--use_mp={}".format(False)
            ]
            p = subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stdout)
            p_list.append(p)
        for p in p_list:
            p.wait()
    else:
        main(args)