predict_system.py 6.69 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
WenmuZhou's avatar
WenmuZhou committed
16

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

LDOUBLEV's avatar
LDOUBLEV committed
21
22
23
24
import cv2
import copy
import numpy as np
import time
LDOUBLEV's avatar
LDOUBLEV committed
25
from PIL import Image
WenmuZhou's avatar
WenmuZhou committed
26
27
28
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
29
import tools.infer.predict_cls as predict_cls
WenmuZhou's avatar
WenmuZhou committed
30
31
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
from ppocr.utils.logging import get_logger
WenmuZhou's avatar
WenmuZhou committed
32
from tools.infer.utility import draw_ocr_box_txt
LDOUBLEV's avatar
LDOUBLEV committed
33

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

LDOUBLEV's avatar
LDOUBLEV committed
36
37
38
39
40

class TextSystem(object):
    def __init__(self, args):
        self.text_detector = predict_det.TextDetector(args)
        self.text_recognizer = predict_rec.TextRecognizer(args)
WenmuZhou's avatar
WenmuZhou committed
41
        self.use_angle_cls = args.use_angle_cls
WenmuZhou's avatar
WenmuZhou committed
42
        self.drop_score = args.drop_score
WenmuZhou's avatar
WenmuZhou committed
43
44
        if self.use_angle_cls:
            self.text_classifier = predict_cls.TextClassifier(args)
LDOUBLEV's avatar
LDOUBLEV committed
45
46

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

    def __call__(self, img):
        ori_im = img.copy()
        dt_boxes, elapse = self.text_detector(img)
WenmuZhou's avatar
WenmuZhou committed
88
89
        logger.info("dt_boxes num : {}, elapse : {}".format(
            len(dt_boxes), elapse))
LDOUBLEV's avatar
LDOUBLEV committed
90
91
92
        if dt_boxes is None:
            return None, None
        img_crop_list = []
93
94
95

        dt_boxes = sorted_boxes(dt_boxes)

LDOUBLEV's avatar
LDOUBLEV committed
96
97
98
99
        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)
WenmuZhou's avatar
WenmuZhou committed
100
101
102
        if self.use_angle_cls:
            img_crop_list, angle_list, elapse = self.text_classifier(
                img_crop_list)
WenmuZhou's avatar
WenmuZhou committed
103
            logger.info("cls num  : {}, elapse : {}".format(
WenmuZhou's avatar
WenmuZhou committed
104
105
                len(img_crop_list), elapse))

LDOUBLEV's avatar
LDOUBLEV committed
106
        rec_res, elapse = self.text_recognizer(img_crop_list)
WenmuZhou's avatar
WenmuZhou committed
107
108
        logger.info("rec_res num  : {}, elapse : {}".format(
            len(rec_res), elapse))
109
        # self.print_draw_crop_rec_res(img_crop_list, rec_res)
WenmuZhou's avatar
WenmuZhou committed
110
111
112
113
114
115
116
        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
117
118


119
120
121
122
def sorted_boxes(dt_boxes):
    """
    Sort text boxes in order from top to bottom, left to right
    args:
tink2123's avatar
tink2123 committed
123
        dt_boxes(array):detected text boxes with shape [4, 2]
124
125
126
127
    return:
        sorted boxes(array) with shape [4, 2]
    """
    num_boxes = dt_boxes.shape[0]
128
    sorted_boxes = sorted(dt_boxes, key=lambda x: (x[0][1], x[0][0]))
129
130
131
    _boxes = list(sorted_boxes)

    for i in range(num_boxes - 1):
WenmuZhou's avatar
WenmuZhou committed
132
133
        if abs(_boxes[i + 1][0][1] - _boxes[i][0][1]) < 10 and \
                (_boxes[i + 1][0][0] < _boxes[i][0][0]):
134
135
136
137
138
139
            tmp = _boxes[i]
            _boxes[i] = _boxes[i + 1]
            _boxes[i + 1] = tmp
    return _boxes


140
def main(args):
LDOUBLEV's avatar
LDOUBLEV committed
141
    image_file_list = get_image_file_list(args.image_dir)
LDOUBLEV's avatar
LDOUBLEV committed
142
    text_sys = TextSystem(args)
LDOUBLEV's avatar
LDOUBLEV committed
143
    is_visualize = True
WenmuZhou's avatar
WenmuZhou committed
144
    font_path = args.vis_font_path
WenmuZhou's avatar
WenmuZhou committed
145
    drop_score = args.drop_score
LDOUBLEV's avatar
LDOUBLEV committed
146
    for image_file in image_file_list:
LDOUBLEV's avatar
LDOUBLEV committed
147
148
149
        img, flag = check_and_read_gif(image_file)
        if not flag:
            img = cv2.imread(image_file)
LDOUBLEV's avatar
LDOUBLEV committed
150
151
152
153
154
155
        if img is None:
            logger.info("error in loading image:{}".format(image_file))
            continue
        starttime = time.time()
        dt_boxes, rec_res = text_sys(img)
        elapse = time.time() - starttime
WenmuZhou's avatar
WenmuZhou committed
156
        logger.info("Predict time of %s: %.3fs" % (image_file, elapse))
LDOUBLEV's avatar
LDOUBLEV committed
157

WenmuZhou's avatar
WenmuZhou committed
158
159
        for text, score in rec_res:
            logger.info("{}, {:.3f}".format(text, score))
LDOUBLEV's avatar
LDOUBLEV committed
160
161
162
163
164
165
166

        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
167
168
169
170
171
172
173
            draw_img = draw_ocr_box_txt(
                image,
                boxes,
                txts,
                scores,
                drop_score=drop_score,
                font_path=font_path)
174
            draw_img_save = "./inference_results/"
LDOUBLEV's avatar
LDOUBLEV committed
175
176
177
178
            if not os.path.exists(draw_img_save):
                os.makedirs(draw_img_save)
            cv2.imwrite(
                os.path.join(draw_img_save, os.path.basename(image_file)),
dyning's avatar
dyning committed
179
                draw_img[:, :, ::-1])
WenmuZhou's avatar
WenmuZhou committed
180
            logger.info("The visualized image saved in {}".format(
181
                os.path.join(draw_img_save, os.path.basename(image_file))))
182
183
184


if __name__ == "__main__":
WenmuZhou's avatar
WenmuZhou committed
185
    main(utility.parse_args())