utility.py 12.4 KB
Newer Older
LDOUBLEV's avatar
LDOUBLEV committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 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.

import argparse
import os, sys
from ppocr.utils.utility import initial_logger
logger = initial_logger()
from paddle.fluid.core import PaddleTensor
from paddle.fluid.core import AnalysisConfig
from paddle.fluid.core import create_paddle_predictor
import cv2
import numpy as np
LDOUBLEV's avatar
LDOUBLEV committed
24
25
import json
from PIL import Image, ImageDraw, ImageFont
26
import math
LDOUBLEV's avatar
LDOUBLEV committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48


def parse_args():
    def str2bool(v):
        return v.lower() in ("true", "t", "1")

    parser = argparse.ArgumentParser()
    #params for prediction engine
    parser.add_argument("--use_gpu", type=str2bool, default=True)
    parser.add_argument("--ir_optim", type=str2bool, default=True)
    parser.add_argument("--use_tensorrt", type=str2bool, default=False)
    parser.add_argument("--gpu_mem", type=int, default=8000)

    #params for text detector
    parser.add_argument("--image_dir", type=str)
    parser.add_argument("--det_algorithm", type=str, default='DB')
    parser.add_argument("--det_model_dir", type=str)
    parser.add_argument("--det_max_side_len", type=float, default=960)

    #DB parmas
    parser.add_argument("--det_db_thresh", type=float, default=0.3)
    parser.add_argument("--det_db_box_thresh", type=float, default=0.5)
49
    parser.add_argument("--det_db_unclip_ratio", type=float, default=2.0)
LDOUBLEV's avatar
LDOUBLEV committed
50
51
52
53
54
55

    #EAST parmas
    parser.add_argument("--det_east_score_thresh", type=float, default=0.8)
    parser.add_argument("--det_east_cover_thresh", type=float, default=0.1)
    parser.add_argument("--det_east_nms_thresh", type=float, default=0.2)

licx's avatar
licx committed
56
57
58
    #SAST parmas
    parser.add_argument("--det_sast_score_thresh", type=float, default=0.5)
    parser.add_argument("--det_sast_nms_thresh", type=float, default=0.2)
59
    parser.add_argument("--det_sast_polygon", type=bool, default=False)
licx's avatar
licx committed
60

LDOUBLEV's avatar
LDOUBLEV committed
61
62
63
    #params for text recognizer
    parser.add_argument("--rec_algorithm", type=str, default='CRNN')
    parser.add_argument("--rec_model_dir", type=str)
tink2123's avatar
fix bug  
tink2123 committed
64
65
    parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320")
    parser.add_argument("--rec_char_type", type=str, default='ch')
66
    parser.add_argument("--rec_batch_num", type=int, default=30)
tink2123's avatar
fix bug  
tink2123 committed
67
    parser.add_argument("--max_text_length", type=int, default=25)
LDOUBLEV's avatar
LDOUBLEV committed
68
69
70
71
    parser.add_argument(
        "--rec_char_dict_path",
        type=str,
        default="./ppocr/utils/ppocr_keys_v1.txt")
tink2123's avatar
tink2123 committed
72
    parser.add_argument("--use_space_char", type=bool, default=True)
dyning's avatar
dyning committed
73
    parser.add_argument("--enable_mkldnn", type=bool, default=False)
littletomatodonkey's avatar
littletomatodonkey committed
74
    parser.add_argument("--use_zero_copy_run", type=bool, default=False)
LDOUBLEV's avatar
LDOUBLEV committed
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
    return parser.parse_args()


def create_predictor(args, mode):
    if mode == "det":
        model_dir = args.det_model_dir
    else:
        model_dir = args.rec_model_dir

    if model_dir is None:
        logger.info("not find {} model file path {}".format(mode, model_dir))
        sys.exit(0)
    model_file_path = model_dir + "/model"
    params_file_path = model_dir + "/params"
    if not os.path.exists(model_file_path):
        logger.info("not find model file path {}".format(model_file_path))
        sys.exit(0)
    if not os.path.exists(params_file_path):
        logger.info("not find params file path {}".format(params_file_path))
        sys.exit(0)

    config = AnalysisConfig(model_file_path, params_file_path)

    if args.use_gpu:
        config.enable_use_gpu(args.gpu_mem, 0)
    else:
        config.disable_gpu()
dyning's avatar
dyning committed
102
103
104
        config.set_cpu_math_library_num_threads(6)
        if args.enable_mkldnn:
            config.enable_mkldnn()
LDOUBLEV's avatar
LDOUBLEV committed
105

tink2123's avatar
tink2123 committed
106
    #config.enable_memory_optim()
LDOUBLEV's avatar
LDOUBLEV committed
107
    config.disable_glog_info()
LDOUBLEV's avatar
LDOUBLEV committed
108

littletomatodonkey's avatar
littletomatodonkey committed
109
110
111
112
113
114
    if args.use_zero_copy_run:
        config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
        config.switch_use_feed_fetch_ops(False)
    else:
        config.switch_use_feed_fetch_ops(True)

LDOUBLEV's avatar
LDOUBLEV committed
115
116
117
118
119
120
121
122
123
124
125
    predictor = create_paddle_predictor(config)
    input_names = predictor.get_input_names()
    input_tensor = predictor.get_input_tensor(input_names[0])
    output_names = predictor.get_output_names()
    output_tensors = []
    for output_name in output_names:
        output_tensor = predictor.get_output_tensor(output_name)
        output_tensors.append(output_tensor)
    return predictor, input_tensor, output_tensors


LDOUBLEV's avatar
LDOUBLEV committed
126
def draw_text_det_res(dt_boxes, img_path):
LDOUBLEV's avatar
LDOUBLEV committed
127
128
129
130
    src_im = cv2.imread(img_path)
    for box in dt_boxes:
        box = np.array(box).astype(np.int32).reshape(-1, 2)
        cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
LDOUBLEV's avatar
LDOUBLEV committed
131
    return src_im
LDOUBLEV's avatar
LDOUBLEV committed
132
133


LDOUBLEV's avatar
LDOUBLEV committed
134
135
def resize_img(img, input_size=600):
    """
LDOUBLEV's avatar
LDOUBLEV committed
136
    resize img and limit the longest side of the image to input_size
LDOUBLEV's avatar
LDOUBLEV committed
137
138
139
140
141
142
143
144
145
146
    """
    img = np.array(img)
    im_shape = img.shape
    im_size_max = np.max(im_shape[0:2])
    im_scale = float(input_size) / float(im_size_max)
    im = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
    return im


def draw_ocr(image, boxes, txts, scores, draw_txt=True, drop_score=0.5):
147
148
149
    """
    Visualize the results of OCR detection and recognition
    args:
LDOUBLEV's avatar
LDOUBLEV committed
150
        image(Image|array): RGB image
151
152
153
154
155
156
157
158
        boxes(list): boxes with shape(N, 4, 2)
        txts(list): the texts
        scores(list): txxs corresponding scores
        draw_txt(bool): whether draw text or not
        drop_score(float): only scores greater than drop_threshold will be visualized
    return(array):
        the visualized img
    """
LDOUBLEV's avatar
LDOUBLEV committed
159
160
    if scores is None:
        scores = [1] * len(boxes)
LDOUBLEV's avatar
LDOUBLEV committed
161
    for (box, score) in zip(boxes, scores):
LDOUBLEV's avatar
LDOUBLEV committed
162
        if score < drop_score or math.isnan(score):
LDOUBLEV's avatar
LDOUBLEV committed
163
            continue
LDOUBLEV's avatar
LDOUBLEV committed
164
        box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64)
LDOUBLEV's avatar
LDOUBLEV committed
165
        image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
LDOUBLEV's avatar
LDOUBLEV committed
166
167

    if draw_txt:
LDOUBLEV's avatar
LDOUBLEV committed
168
        img = np.array(resize_img(image, input_size=600))
169
170
171
        txt_img = text_visual(
            txts, scores, img_h=img.shape[0], img_w=600, threshold=drop_score)
        img = np.concatenate([np.array(img), np.array(txt_img)], axis=1)
LDOUBLEV's avatar
LDOUBLEV committed
172
173
        return img
    return image
174
175


176
177
178
179
def draw_ocr_box_txt(image, boxes, txts):
    h, w = image.height, image.width
    img_left = image.copy()
    img_right = Image.new('RGB', (w, h), (255, 255, 255))
180
181

    import random
LDOUBLEV's avatar
LDOUBLEV committed
182

183
184
185
    random.seed(0)
    draw_left = ImageDraw.Draw(img_left)
    draw_right = ImageDraw.Draw(img_right)
186
    for (box, txt) in zip(boxes, txts):
tink2123's avatar
tink2123 committed
187
188
        color = (random.randint(0, 255), random.randint(0, 255),
                 random.randint(0, 255))
189
        draw_left.polygon(box, fill=color)
tink2123's avatar
tink2123 committed
190
191
192
193
194
195
196
197
198
199
        draw_right.polygon(
            [
                box[0][0], box[0][1], box[1][0], box[1][1], box[2][0],
                box[2][1], box[3][0], box[3][1]
            ],
            outline=color)
        box_height = math.sqrt((box[0][0] - box[3][0])**2 + (box[0][1] - box[3][
            1])**2)
        box_width = math.sqrt((box[0][0] - box[1][0])**2 + (box[0][1] - box[1][
            1])**2)
200
201
        if box_height > 2 * box_width:
            font_size = max(int(box_width * 0.9), 10)
tink2123's avatar
tink2123 committed
202
203
            font = ImageFont.truetype(
                "./doc/simfang.ttf", font_size, encoding="utf-8")
204
205
206
            cur_y = box[0][1]
            for c in txt:
                char_size = font.getsize(c)
tink2123's avatar
tink2123 committed
207
208
                draw_right.text(
                    (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
209
210
211
                cur_y += char_size[1]
        else:
            font_size = max(int(box_height * 0.8), 10)
tink2123's avatar
tink2123 committed
212
213
214
215
            font = ImageFont.truetype(
                "./doc/simfang.ttf", font_size, encoding="utf-8")
            draw_right.text(
                [box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
216
217
218
219
    img_left = Image.blend(image, img_left, 0.5)
    img_show = Image.new('RGB', (w * 2, h), (255, 255, 255))
    img_show.paste(img_left, (0, 0, w, h))
    img_show.paste(img_right, (w, 0, w * 2, h))
220
221
222
    return np.array(img_show)


223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
def str_count(s):
    """
    Count the number of Chinese characters,
    a single English character and a single number
    equal to half the length of Chinese characters.

    args:
        s(string): the input of string
    return(int):
        the number of Chinese characters
    """
    import string
    count_zh = count_pu = 0
    s_len = len(s)
    en_dg_count = 0
    for c in s:
        if c in string.ascii_letters or c.isdigit() or c.isspace():
            en_dg_count += 1
        elif c.isalpha():
            count_zh += 1
        else:
            count_pu += 1
    return s_len - math.ceil(en_dg_count / 2)


def text_visual(texts, scores, img_h=400, img_w=600, threshold=0.):
    """
    create new blank img and draw txt on it
    args:
        texts(list): the text will be draw
        scores(list|None): corresponding score of each txt
        img_h(int): the height of blank img
        img_w(int): the width of blank img
    return(array):

    """
    if scores is not None:
        assert len(texts) == len(
            scores), "The number of txts and corresponding scores must match"

    def create_blank_img():
        blank_img = np.ones(shape=[img_h, img_w], dtype=np.int8) * 255
        blank_img[:, img_w - 1:] = 0
LDOUBLEV's avatar
LDOUBLEV committed
266
267
        blank_img = Image.fromarray(blank_img).convert("RGB")
        draw_txt = ImageDraw.Draw(blank_img)
268
        return blank_img, draw_txt
LDOUBLEV's avatar
LDOUBLEV committed
269

270
271
272
273
    blank_img, draw_txt = create_blank_img()

    font_size = 20
    txt_color = (0, 0, 0)
274
    font = ImageFont.truetype("./doc/simfang.ttf", font_size, encoding="utf-8")
275
276
277

    gap = font_size + 5
    txt_img_list = []
LDOUBLEV's avatar
LDOUBLEV committed
278
    count, index = 1, 0
279
280
    for idx, txt in enumerate(texts):
        index += 1
LDOUBLEV's avatar
LDOUBLEV committed
281
        if scores[idx] < threshold or math.isnan(scores[idx]):
282
283
284
285
286
287
288
289
290
291
292
            index -= 1
            continue
        first_line = True
        while str_count(txt) >= img_w // font_size - 4:
            tmp = txt
            txt = tmp[:img_w // font_size - 4]
            if first_line:
                new_txt = str(index) + ': ' + txt
                first_line = False
            else:
                new_txt = '    ' + txt
LDOUBLEV's avatar
LDOUBLEV committed
293
            draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
294
295
296
297
298
            txt = tmp[img_w // font_size - 4:]
            if count >= img_h // gap - 1:
                txt_img_list.append(np.array(blank_img))
                blank_img, draw_txt = create_blank_img()
                count = 0
LDOUBLEV's avatar
LDOUBLEV committed
299
            count += 1
300
301
302
        if first_line:
            new_txt = str(index) + ': ' + txt + '   ' + '%.3f' % (scores[idx])
        else:
LDOUBLEV's avatar
LDOUBLEV committed
303
            new_txt = "  " + txt + "  " + '%.3f' % (scores[idx])
LDOUBLEV's avatar
LDOUBLEV committed
304
        draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
305
        # whether add new blank img or not
LDOUBLEV's avatar
LDOUBLEV committed
306
        if count >= img_h // gap - 1 and idx + 1 < len(texts):
307
308
309
            txt_img_list.append(np.array(blank_img))
            blank_img, draw_txt = create_blank_img()
            count = 0
LDOUBLEV's avatar
LDOUBLEV committed
310
        count += 1
311
312
313
314
315
316
    txt_img_list.append(np.array(blank_img))
    if len(txt_img_list) == 1:
        blank_img = np.array(txt_img_list[0])
    else:
        blank_img = np.concatenate(txt_img_list, axis=1)
    return np.array(blank_img)
LDOUBLEV's avatar
LDOUBLEV committed
317
318


dyning's avatar
dyning committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def base64_to_cv2(b64str):
    import base64
    data = base64.b64decode(b64str.encode('utf8'))
    data = np.fromstring(data, np.uint8)
    data = cv2.imdecode(data, cv2.IMREAD_COLOR)
    return data


def draw_boxes(image, boxes, scores=None, drop_score=0.5):
    if scores is None:
        scores = [1] * len(boxes)
    for (box, score) in zip(boxes, scores):
        if score < drop_score:
            continue
        box = np.reshape(np.array(box), [-1, 1, 2]).astype(np.int64)
        image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
    return image


LDOUBLEV's avatar
LDOUBLEV committed
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
if __name__ == '__main__':
    test_img = "./doc/test_v2"
    predict_txt = "./doc/predict.txt"
    f = open(predict_txt, 'r')
    data = f.readlines()
    img_path, anno = data[0].strip().split('\t')
    img_name = os.path.basename(img_path)
    img_path = os.path.join(test_img, img_name)
    image = Image.open(img_path)

    data = json.loads(anno)
    boxes, txts, scores = [], [], []
    for dic in data:
        boxes.append(dic['points'])
        txts.append(dic['transcription'])
        scores.append(round(dic['scores'], 3))

    new_img = draw_ocr(image, boxes, txts, scores, draw_txt=True)

MissPenguin's avatar
MissPenguin committed
357
    cv2.imwrite(img_name, new_img)