predict_system.py 8.04 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
from tools.infer.utility import draw_ocr_box_txt
WenmuZhou's avatar
WenmuZhou committed
37
38
logger = get_logger()

LDOUBLEV's avatar
LDOUBLEV committed
39
40
41

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

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

    def get_rotate_crop_image(self, img, points):
53
        '''
LDOUBLEV's avatar
LDOUBLEV committed
54
55
56
57
58
59
60
61
        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
62
        '''
LDOUBLEV's avatar
LDOUBLEV committed
63
64
65
66
67
68
69
70
71
        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],
72
73
                              [img_crop_width, img_crop_height],
                              [0, img_crop_height]])
LDOUBLEV's avatar
LDOUBLEV committed
74
        M = cv2.getPerspectiveTransform(points, pts_std)
LDOUBLEV's avatar
LDOUBLEV committed
75
76
77
78
79
        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
80
81
82
83
84
85
86
87
88
        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
89
            logger.info(bno, rec_res[bno])
LDOUBLEV's avatar
LDOUBLEV committed
90

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

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

        dt_boxes = sorted_boxes(dt_boxes)

LDOUBLEV's avatar
LDOUBLEV committed
103
104
105
106
        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)
107
        if self.use_angle_cls and cls:
WenmuZhou's avatar
WenmuZhou committed
108
109
            img_crop_list, angle_list, elapse = self.text_classifier(
                img_crop_list)
WenmuZhou's avatar
WenmuZhou committed
110
            logger.debug("cls num  : {}, elapse : {}".format(
WenmuZhou's avatar
WenmuZhou committed
111
112
                len(img_crop_list), elapse))

LDOUBLEV's avatar
LDOUBLEV committed
113
        rec_res, elapse = self.text_recognizer(img_crop_list)
WenmuZhou's avatar
WenmuZhou committed
114
        logger.debug("rec_res num  : {}, elapse : {}".format(
WenmuZhou's avatar
WenmuZhou committed
115
            len(rec_res), elapse))
116
        # self.print_draw_crop_rec_res(img_crop_list, rec_res)
WenmuZhou's avatar
WenmuZhou committed
117
118
119
120
121
122
123
        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
124
125


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

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


147
def main(args):
LDOUBLEV's avatar
LDOUBLEV committed
148
    image_file_list = get_image_file_list(args.image_dir)
LDOUBLEV's avatar
LDOUBLEV committed
149
    image_file_list = image_file_list[args.process_id::args.total_process_num]
LDOUBLEV's avatar
LDOUBLEV committed
150
    text_sys = TextSystem(args)
LDOUBLEV's avatar
LDOUBLEV committed
151
    is_visualize = True
WenmuZhou's avatar
WenmuZhou committed
152
    font_path = args.vis_font_path
WenmuZhou's avatar
WenmuZhou committed
153
    drop_score = args.drop_score
Double_V's avatar
Double_V committed
154

LDOUBLEV's avatar
LDOUBLEV committed
155
156
157
158
159
    # warm up 10 times
    if args.warmup:
        img = np.random.uniform(0, 255, [640, 640, 3]).astype(np.uint8)
        for i in range(10):
            res = text_sys(img)
LDOUBLEV's avatar
LDOUBLEV committed
160

LDOUBLEV's avatar
LDOUBLEV committed
161
162
163
164
165
    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
166

LDOUBLEV's avatar
LDOUBLEV committed
167
168
169
        img, flag = check_and_read_gif(image_file)
        if not flag:
            img = cv2.imread(image_file)
LDOUBLEV's avatar
LDOUBLEV committed
170
        if img is None:
171
            logger.info("error in loading image:{}".format(image_file))
LDOUBLEV's avatar
LDOUBLEV committed
172
173
174
175
            continue
        starttime = time.time()
        dt_boxes, rec_res = text_sys(img)
        elapse = time.time() - starttime
LDOUBLEV's avatar
LDOUBLEV committed
176
        total_time += elapse
LDOUBLEV's avatar
LDOUBLEV committed
177

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

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

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

LDOUBLEV's avatar
LDOUBLEV committed
210
211
212
213
    img_num = text_sys.text_detector.det_times.img_num


if __name__ == "__main__":
LDOUBLEV's avatar
LDOUBLEV committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
    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)