predict_det.py 6.06 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
16
17
18
import os
import sys
__dir__ = os.path.dirname(__file__)
sys.path.append(__dir__)
sys.path.append(os.path.join(__dir__, '../..'))
LDOUBLEV's avatar
LDOUBLEV committed
19

LDOUBLEV's avatar
LDOUBLEV committed
20
import tools.infer.utility as utility
LDOUBLEV's avatar
LDOUBLEV committed
21
22
from ppocr.utils.utility import initial_logger
logger = initial_logger()
LDOUBLEV's avatar
LDOUBLEV committed
23
from ppocr.utils.utility import get_image_file_list, check_and_read_gif
LDOUBLEV's avatar
LDOUBLEV committed
24
25
26
27
28
29
30
31
32
import cv2
from ppocr.data.det.east_process import EASTProcessTest
from ppocr.data.det.db_process import DBProcessTest
from ppocr.postprocess.db_postprocess import DBPostProcess
from ppocr.postprocess.east_postprocess import EASTPostPocess
import copy
import numpy as np
import math
import time
33
import sys
LDOUBLEV's avatar
LDOUBLEV committed
34
35
36
37
38
39
40
41
42
43
44
45
46


class TextDetector(object):
    def __init__(self, args):
        max_side_len = args.det_max_side_len
        self.det_algorithm = args.det_algorithm
        preprocess_params = {'max_side_len': max_side_len}
        postprocess_params = {}
        if self.det_algorithm == "DB":
            self.preprocess_op = DBProcessTest(preprocess_params)
            postprocess_params["thresh"] = args.det_db_thresh
            postprocess_params["box_thresh"] = args.det_db_box_thresh
            postprocess_params["max_candidates"] = 1000
47
            postprocess_params["unclip_ratio"] = args.det_db_unclip_ratio
LDOUBLEV's avatar
LDOUBLEV committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
            self.postprocess_op = DBPostProcess(postprocess_params)
        elif self.det_algorithm == "EAST":
            self.preprocess_op = EASTProcessTest(preprocess_params)
            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
            self.postprocess_op = EASTPostPocess(postprocess_params)
        else:
            logger.info("unknown det_algorithm:{}".format(self.det_algorithm))
            sys.exit(0)

        self.predictor, self.input_tensor, self.output_tensors =\
            utility.create_predictor(args, mode="det")

    def order_points_clockwise(self, pts):
63
64
        """
        reference from: https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py
LDOUBLEV's avatar
LDOUBLEV committed
65
        # sort the points based on their x-coordinates
66
        """
LDOUBLEV's avatar
LDOUBLEV committed
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
        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
86
87
88
89
    def clip_det_res(self, points, img_height, img_width):
        for pno in range(4):
            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
90
91
92
93
94
95
96
        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
97
            box = self.clip_det_res(box, img_height, img_width)
LDOUBLEV's avatar
LDOUBLEV committed
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
            rect_width = int(np.linalg.norm(box[0] - box[1]))
            rect_height = int(np.linalg.norm(box[0] - box[3]))
            if rect_width <= 10 or rect_height <= 10:
                continue
            dt_boxes_new.append(box)
        dt_boxes = np.array(dt_boxes_new)
        return dt_boxes

    def __call__(self, img):
        ori_im = img.copy()
        im, ratio_list = self.preprocess_op(img)
        if im is None:
            return None, 0
        im = im.copy()
        starttime = time.time()
        self.input_tensor.copy_from_cpu(im)
        self.predictor.zero_copy_run()
        outputs = []
        for output_tensor in self.output_tensors:
            output = output_tensor.copy_to_cpu()
            outputs.append(output)
        outs_dict = {}
        if self.det_algorithm == "EAST":
LDOUBLEV's avatar
LDOUBLEV committed
121
122
            outs_dict['f_geo'] = outputs[0]
            outs_dict['f_score'] = outputs[1]
LDOUBLEV's avatar
LDOUBLEV committed
123
        else:
124
            outs_dict['maps'] = outputs[0]
LDOUBLEV's avatar
LDOUBLEV committed
125
126
127
128
129
130
131
132
133
        dt_boxes_list = self.postprocess_op(outs_dict, [ratio_list])
        dt_boxes = dt_boxes_list[0]
        dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape)
        elapse = time.time() - starttime
        return dt_boxes, elapse


if __name__ == "__main__":
    args = utility.parse_args()
LDOUBLEV's avatar
LDOUBLEV committed
134
    image_file_list = get_image_file_list(args.image_dir)
LDOUBLEV's avatar
LDOUBLEV committed
135
136
137
    text_detector = TextDetector(args)
    count = 0
    total_time = 0
littletomatodonkey's avatar
littletomatodonkey committed
138
139
140
    draw_img_save = "./inference_results"
    if not os.path.exists(draw_img_save):
        os.makedirs(draw_img_save)
LDOUBLEV's avatar
LDOUBLEV committed
141
    for image_file in image_file_list:
LDOUBLEV's avatar
LDOUBLEV committed
142
143
144
        img, flag = check_and_read_gif(image_file)
        if not flag:
            img = cv2.imread(image_file)
LDOUBLEV's avatar
LDOUBLEV committed
145
146
147
148
149
150
151
152
        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
        print("Predict time of %s:" % image_file, elapse)
dyning's avatar
dyning committed
153
154
        src_im = utility.draw_text_det_res(dt_boxes, image_file)
        img_name_pure = image_file.split("/")[-1]
littletomatodonkey's avatar
littletomatodonkey committed
155
156
        cv2.imwrite(
            os.path.join(draw_img_save, "det_res_%s" % img_name_pure), src_im)
157
158
    if count > 1:
        print("Avg Time:", total_time / (count - 1))