predict_det.py 8.48 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
23
24
import cv2
import numpy as np
import time
import sys
WenmuZhou's avatar
WenmuZhou committed
25
import paddle
26

LDOUBLEV's avatar
LDOUBLEV committed
27
import tools.infer.utility as utility
WenmuZhou's avatar
WenmuZhou committed
28
from ppocr.utils.logging import get_logger
LDOUBLEV's avatar
LDOUBLEV committed
29
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
WenmuZhou's avatar
WenmuZhou committed
30
31
from ppocr.data import create_operators, transform
from ppocr.postprocess import build_post_process
LDOUBLEV's avatar
LDOUBLEV committed
32

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

LDOUBLEV's avatar
LDOUBLEV committed
35
36
37

class TextDetector(object):
    def __init__(self, args):
LDOUBLEV's avatar
LDOUBLEV committed
38
        self.args = args
LDOUBLEV's avatar
LDOUBLEV committed
39
        self.det_algorithm = args.det_algorithm
MissPenguin's avatar
MissPenguin committed
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
        pre_process_list = [{
            'DetResizeForTest': {
                'limit_side_len': args.det_limit_side_len,
                'limit_type': args.det_limit_type
            }
        }, {
            'NormalizeImage': {
                'std': [0.229, 0.224, 0.225],
                'mean': [0.485, 0.456, 0.406],
                'scale': '1./255.',
                'order': 'hwc'
            }
        }, {
            'ToCHWImage': None
        }, {
            'KeepKeys': {
                'keep_keys': ['image', 'shape']
            }
        }]
LDOUBLEV's avatar
LDOUBLEV committed
59
60
        postprocess_params = {}
        if self.det_algorithm == "DB":
WenmuZhou's avatar
WenmuZhou committed
61
            postprocess_params['name'] = 'DBPostProcess'
LDOUBLEV's avatar
LDOUBLEV committed
62
63
64
            postprocess_params["thresh"] = args.det_db_thresh
            postprocess_params["box_thresh"] = args.det_db_box_thresh
            postprocess_params["max_candidates"] = 1000
65
            postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio
MissPenguin's avatar
MissPenguin committed
66
            postprocess_params["use_dilation"] = True
MissPenguin's avatar
MissPenguin committed
67
        elif self.det_algorithm == "EAST":
WenmuZhou's avatar
WenmuZhou committed
68
            postprocess_params['name'] = 'EASTPostProcess'
MissPenguin's avatar
MissPenguin committed
69
70
71
72
            postprocess_params["score_thresh"] = args.det_east_score_thresh
            postprocess_params["cover_thresh"] = args.det_east_cover_thresh
            postprocess_params["nms_thresh"] = args.det_east_nms_thresh
        elif self.det_algorithm == "SAST":
MissPenguin's avatar
MissPenguin committed
73
            pre_process_list[0] = {
WenmuZhou's avatar
WenmuZhou committed
74
75
76
                'DetResizeForTest': {
                    'resize_long': args.det_limit_side_len
                }
MissPenguin's avatar
MissPenguin committed
77
            }
WenmuZhou's avatar
WenmuZhou committed
78
            postprocess_params['name'] = 'SASTPostProcess'
MissPenguin's avatar
MissPenguin committed
79
80
81
82
83
84
85
86
87
88
89
            postprocess_params["score_thresh"] = args.det_sast_score_thresh
            postprocess_params["nms_thresh"] = args.det_sast_nms_thresh
            self.det_sast_polygon = args.det_sast_polygon
            if self.det_sast_polygon:
                postprocess_params["sample_pts_num"] = 6
                postprocess_params["expand_scale"] = 1.2
                postprocess_params["shrink_ratio_of_width"] = 0.2
            else:
                postprocess_params["sample_pts_num"] = 2
                postprocess_params["expand_scale"] = 1.0
                postprocess_params["shrink_ratio_of_width"] = 0.3
LDOUBLEV's avatar
LDOUBLEV committed
90
91
92
93
        else:
            logger.info("unknown det_algorithm:{}".format(self.det_algorithm))
            sys.exit(0)

WenmuZhou's avatar
WenmuZhou committed
94
95
        self.preprocess_op = create_operators(pre_process_list)
        self.postprocess_op = build_post_process(postprocess_params)
96
97
98
        self.predictor, self.input_tensor, self.output_tensors = utility.create_predictor(
            args, 'det', logger)  # paddle.jit.load(args.det_model_dir)
        # self.predictor.eval()
LDOUBLEV's avatar
LDOUBLEV committed
99
100

    def order_points_clockwise(self, pts):
101
102
        """
        reference from: https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
LDOUBLEV's avatar
LDOUBLEV committed
103
        # sort the points based on their x-coordinates
104
        """
LDOUBLEV's avatar
LDOUBLEV committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
        xSorted = pts[np.argsort(pts[:, 0]), :]

        # grab the left-most and right-most points from the sorted
        # x-roodinate points
        leftMost = xSorted[:2, :]
        rightMost = xSorted[2:, :]

        # now, sort the left-most coordinates according to their
        # y-coordinates so we can grab the top-left and bottom-left
        # points, respectively
        leftMost = leftMost[np.argsort(leftMost[:, 1]), :]
        (tl, bl) = leftMost

        rightMost = rightMost[np.argsort(rightMost[:, 1]), :]
        (tr, br) = rightMost

        rect = np.array([tl, tr, br, bl], dtype="float32")
        return rect

dyning's avatar
dyning committed
124
    def clip_det_res(self, points, img_height, img_width):
125
        for pno in range(points.shape[0]):
dyning's avatar
dyning committed
126
127
            points[pno, 0] = int(min(max(points[pno, 0], 0), img_width - 1))
            points[pno, 1] = int(min(max(points[pno, 1], 0), img_height - 1))
LDOUBLEV's avatar
LDOUBLEV committed
128
129
130
131
132
133
134
        return points

    def filter_tag_det_res(self, dt_boxes, image_shape):
        img_height, img_width = image_shape[0:2]
        dt_boxes_new = []
        for box in dt_boxes:
            box = self.order_points_clockwise(box)
dyning's avatar
dyning committed
135
            box = self.clip_det_res(box, img_height, img_width)
LDOUBLEV's avatar
LDOUBLEV committed
136
137
            rect_width = int(np.linalg.norm(box[0] - box[1]))
            rect_height = int(np.linalg.norm(box[0] - box[3]))
MissPenguin's avatar
MissPenguin committed
138
            if rect_width <= 3 or rect_height <= 3:
LDOUBLEV's avatar
LDOUBLEV committed
139
140
141
142
143
                continue
            dt_boxes_new.append(box)
        dt_boxes = np.array(dt_boxes_new)
        return dt_boxes

144
145
146
147
148
149
150
151
    def filter_tag_det_res_only_clip(self, dt_boxes, image_shape):
        img_height, img_width = image_shape[0:2]
        dt_boxes_new = []
        for box in dt_boxes:
            box = self.clip_det_res(box, img_height, img_width)
            dt_boxes_new.append(box)
        dt_boxes = np.array(dt_boxes_new)
        return dt_boxes
152

LDOUBLEV's avatar
LDOUBLEV committed
153
154
    def __call__(self, img):
        ori_im = img.copy()
WenmuZhou's avatar
WenmuZhou committed
155
156
157
158
        data = {'image': img}
        data = transform(data, self.preprocess_op)
        img, shape_list = data
        if img is None:
LDOUBLEV's avatar
LDOUBLEV committed
159
            return None, 0
WenmuZhou's avatar
WenmuZhou committed
160
161
        img = np.expand_dims(img, axis=0)
        shape_list = np.expand_dims(shape_list, axis=0)
162
        img = img.copy()
LDOUBLEV's avatar
LDOUBLEV committed
163
        starttime = time.time()
164

WenmuZhou's avatar
WenmuZhou committed
165
166
        self.input_tensor.copy_from_cpu(img)
        self.predictor.run()
167
168
169
170
171
        outputs = []
        for output_tensor in self.output_tensors:
            output = output_tensor.copy_to_cpu()
            outputs.append(output)

MissPenguin's avatar
MissPenguin committed
172
173
174
175
176
177
178
179
180
        preds = {}
        if self.det_algorithm == "EAST":
            preds['f_geo'] = outputs[0]
            preds['f_score'] = outputs[1]
        elif self.det_algorithm == 'SAST':
            preds['f_border'] = outputs[0]
            preds['f_score'] = outputs[1]
            preds['f_tco'] = outputs[2]
            preds['f_tvo'] = outputs[3]
WenmuZhou's avatar
WenmuZhou committed
181
        elif self.det_algorithm == 'DB':
WenmuZhou's avatar
WenmuZhou committed
182
            preds['maps'] = outputs[0]
WenmuZhou's avatar
WenmuZhou committed
183
184
        else:
            raise NotImplementedError
MissPenguin's avatar
MissPenguin committed
185

WenmuZhou's avatar
WenmuZhou committed
186
187
        post_result = self.postprocess_op(preds, shape_list)
        dt_boxes = post_result[0]['points']
MissPenguin's avatar
MissPenguin committed
188
189
190
191
        if self.det_algorithm == "SAST" and self.det_sast_polygon:
            dt_boxes = self.filter_tag_det_res_only_clip(dt_boxes, ori_im.shape)
        else:
            dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
LDOUBLEV's avatar
LDOUBLEV committed
192
193
194
195
196
197
        elapse = time.time() - starttime
        return dt_boxes, elapse


if __name__ == "__main__":
    args = utility.parse_args()
LDOUBLEV's avatar
LDOUBLEV committed
198
    image_file_list = get_image_file_list(args.image_dir)
LDOUBLEV's avatar
LDOUBLEV committed
199
200
201
    text_detector = TextDetector(args)
    count = 0
    total_time = 0
littletomatodonkey's avatar
littletomatodonkey committed
202
203
204
    draw_img_save = "./inference_results"
    if not os.path.exists(draw_img_save):
        os.makedirs(draw_img_save)
LDOUBLEV's avatar
LDOUBLEV committed
205
    for image_file in image_file_list:
LDOUBLEV's avatar
LDOUBLEV committed
206
207
208
        img, flag = check_and_read_gif(image_file)
        if not flag:
            img = cv2.imread(image_file)
LDOUBLEV's avatar
LDOUBLEV committed
209
210
211
212
213
214
215
        if img is None:
            logger.info("error in loading image:{}".format(image_file))
            continue
        dt_boxes, elapse = text_detector(img)
        if count > 0:
            total_time += elapse
        count += 1
WenmuZhou's avatar
WenmuZhou committed
216
        logger.info("Predict time of {}: {}".format(image_file, elapse))
dyning's avatar
dyning committed
217
        src_im = utility.draw_text_det_res(dt_boxes, image_file)
WenmuZhou's avatar
WenmuZhou committed
218
        img_name_pure = os.path.split(image_file)[-1]
WenmuZhou's avatar
WenmuZhou committed
219
220
        img_path = os.path.join(draw_img_save,
                                "det_res_{}".format(img_name_pure))
WenmuZhou's avatar
WenmuZhou committed
221
        cv2.imwrite(img_path, src_im)
WenmuZhou's avatar
WenmuZhou committed
222
        logger.info("The visualized image saved in {}".format(img_path))
223
    if count > 1:
224
        logger.info("Avg Time: {}".format(total_time / (count - 1)))