predict_det.py 8.53 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
littletomatodonkey's avatar
littletomatodonkey committed
40
        self.use_zero_copy_run = args.use_zero_copy_run
MissPenguin's avatar
MissPenguin committed
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
        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
60
61
        postprocess_params = {}
        if self.det_algorithm == "DB":
WenmuZhou's avatar
WenmuZhou committed
62
            postprocess_params['name'] = 'DBPostProcess'
LDOUBLEV's avatar
LDOUBLEV committed
63
64
65
            postprocess_params["thresh"] = args.det_db_thresh
            postprocess_params["box_thresh"] = args.det_db_box_thresh
            postprocess_params["max_candidates"] = 1000
66
            postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio
MissPenguin's avatar
MissPenguin committed
67
            postprocess_params["use_dilation"] = True
MissPenguin's avatar
MissPenguin committed
68
        elif self.det_algorithm == "EAST":
WenmuZhou's avatar
WenmuZhou committed
69
            postprocess_params['name'] = 'EASTPostProcess'
MissPenguin's avatar
MissPenguin committed
70
71
72
73
            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":
WenmuZhou's avatar
WenmuZhou committed
74
            postprocess_params['name'] = 'SASTPostProcess'
MissPenguin's avatar
MissPenguin committed
75
76
77
78
79
80
81
82
83
84
85
            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
86
87
88
89
        else:
            logger.info("unknown det_algorithm:{}".format(self.det_algorithm))
            sys.exit(0)

WenmuZhou's avatar
WenmuZhou committed
90
91
        self.preprocess_op = create_operators(pre_process_list)
        self.postprocess_op = build_post_process(postprocess_params)
92
93
94
        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
95
96

    def order_points_clockwise(self, pts):
97
98
        """
        reference from: https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
LDOUBLEV's avatar
LDOUBLEV committed
99
        # sort the points based on their x-coordinates
100
        """
LDOUBLEV's avatar
LDOUBLEV committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
        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
120
    def clip_det_res(self, points, img_height, img_width):
121
        for pno in range(points.shape[0]):
dyning's avatar
dyning committed
122
123
            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
124
125
126
127
128
129
130
        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
131
            box = self.clip_det_res(box, img_height, img_width)
LDOUBLEV's avatar
LDOUBLEV committed
132
133
            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
134
            if rect_width <= 3 or rect_height <= 3:
LDOUBLEV's avatar
LDOUBLEV committed
135
136
137
138
139
                continue
            dt_boxes_new.append(box)
        dt_boxes = np.array(dt_boxes_new)
        return dt_boxes

140
141
142
143
144
145
146
147
    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
148

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

161
162
163
164
165
166
167
168
169
170
171
        if self.use_zero_copy_run:
            self.input_tensor.copy_from_cpu(img)
            self.predictor.zero_copy_run()
        else:
            im = paddle.fluid.core.PaddleTensor(img)
            self.predictor.run([im])
        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)))