utility.py 22.2 KB
Newer Older
L
LDOUBLEV 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
# 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
W
WenmuZhou 已提交
16
import os
W
WenmuZhou 已提交
17
import sys
L
LDOUBLEV 已提交
18 19
import cv2
import numpy as np
L
LDOUBLEV 已提交
20 21
import json
from PIL import Image, ImageDraw, ImageFont
22
import math
W
WenmuZhou 已提交
23
from paddle import inference
L
LDOUBLEV 已提交
24 25
import time
from ppocr.utils.logging import get_logger
W
WenmuZhou 已提交
26

L
LDOUBLEV 已提交
27

28 29
def str2bool(v):
    return v.lower() in ("true", "t", "1")
L
LDOUBLEV 已提交
30 31


W
WenmuZhou 已提交
32
def init_args():
L
LDOUBLEV 已提交
33
    parser = argparse.ArgumentParser()
W
WenmuZhou 已提交
34
    # params for prediction engine
L
LDOUBLEV 已提交
35 36 37
    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)
L
LDOUBLEV 已提交
38
    parser.add_argument("--min_subgraph_size", type=int, default=15)
L
LDOUBLEV 已提交
39
    parser.add_argument("--precision", type=str, default="fp32")
L
LDOUBLEV 已提交
40
    parser.add_argument("--gpu_mem", type=int, default=500)
L
LDOUBLEV 已提交
41

W
WenmuZhou 已提交
42
    # params for text detector
L
LDOUBLEV 已提交
43 44 45
    parser.add_argument("--image_dir", type=str)
    parser.add_argument("--det_algorithm", type=str, default='DB')
    parser.add_argument("--det_model_dir", type=str)
W
WenmuZhou 已提交
46 47
    parser.add_argument("--det_limit_side_len", type=float, default=960)
    parser.add_argument("--det_limit_type", type=str, default='max')
L
LDOUBLEV 已提交
48

W
WenmuZhou 已提交
49
    # DB parmas
L
LDOUBLEV 已提交
50
    parser.add_argument("--det_db_thresh", type=float, default=0.3)
L
LDOUBLEV 已提交
51 52
    parser.add_argument("--det_db_box_thresh", type=float, default=0.6)
    parser.add_argument("--det_db_unclip_ratio", type=float, default=1.5)
L
LDOUBLEV 已提交
53
    parser.add_argument("--max_batch_size", type=int, default=10)
littletomatodonkey's avatar
littletomatodonkey 已提交
54
    parser.add_argument("--use_dilation", type=str2bool, default=False)
littletomatodonkey's avatar
littletomatodonkey 已提交
55
    parser.add_argument("--det_db_score_mode", type=str, default="fast")
W
WenmuZhou 已提交
56
    # EAST parmas
L
LDOUBLEV 已提交
57 58 59 60
    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)

W
WenmuZhou 已提交
61
    # SAST parmas
L
licx 已提交
62 63
    parser.add_argument("--det_sast_score_thresh", type=float, default=0.5)
    parser.add_argument("--det_sast_nms_thresh", type=float, default=0.2)
littletomatodonkey's avatar
littletomatodonkey 已提交
64
    parser.add_argument("--det_sast_polygon", type=str2bool, default=False)
L
licx 已提交
65

W
WenmuZhou 已提交
66
    # params for text recognizer
L
LDOUBLEV 已提交
67 68
    parser.add_argument("--rec_algorithm", type=str, default='CRNN')
    parser.add_argument("--rec_model_dir", type=str)
T
fix bug  
tink2123 已提交
69 70
    parser.add_argument("--rec_image_shape", type=str, default="3, 32, 320")
    parser.add_argument("--rec_char_type", type=str, default='ch')
L
LDOUBLEV 已提交
71
    parser.add_argument("--rec_batch_num", type=int, default=6)
T
fix bug  
tink2123 已提交
72
    parser.add_argument("--max_text_length", type=int, default=25)
L
LDOUBLEV 已提交
73 74 75 76
    parser.add_argument(
        "--rec_char_dict_path",
        type=str,
        default="./ppocr/utils/ppocr_keys_v1.txt")
W
WenmuZhou 已提交
77 78
    parser.add_argument("--use_space_char", type=str2bool, default=True)
    parser.add_argument(
T
tink2123 已提交
79
        "--vis_font_path", type=str, default="./doc/fonts/simfang.ttf")
W
WenmuZhou 已提交
80
    parser.add_argument("--drop_score", type=float, default=0.5)
W
WenmuZhou 已提交
81

J
Jethong 已提交
82 83 84 85 86 87 88 89 90
    # params for e2e
    parser.add_argument("--e2e_algorithm", type=str, default='PGNet')
    parser.add_argument("--e2e_model_dir", type=str)
    parser.add_argument("--e2e_limit_side_len", type=float, default=768)
    parser.add_argument("--e2e_limit_type", type=str, default='max')

    # PGNet parmas
    parser.add_argument("--e2e_pgnet_score_thresh", type=float, default=0.5)
    parser.add_argument(
J
Jethong 已提交
91
        "--e2e_char_dict_path", type=str, default="./ppocr/utils/ic15_dict.txt")
J
Jethong 已提交
92
    parser.add_argument("--e2e_pgnet_valid_set", type=str, default='totaltext')
littletomatodonkey's avatar
littletomatodonkey 已提交
93
    parser.add_argument("--e2e_pgnet_polygon", type=str2bool, default=True)
J
Jethong 已提交
94
    parser.add_argument("--e2e_pgnet_mode", type=str, default='fast')
J
Jethong 已提交
95

W
WenmuZhou 已提交
96 97 98 99 100
    # params for text classifier
    parser.add_argument("--use_angle_cls", type=str2bool, default=False)
    parser.add_argument("--cls_model_dir", type=str)
    parser.add_argument("--cls_image_shape", type=str, default="3, 48, 192")
    parser.add_argument("--label_list", type=list, default=['0', '180'])
L
LDOUBLEV 已提交
101
    parser.add_argument("--cls_batch_num", type=int, default=6)
W
WenmuZhou 已提交
102 103 104
    parser.add_argument("--cls_thresh", type=float, default=0.9)

    parser.add_argument("--enable_mkldnn", type=str2bool, default=False)
L
LDOUBLEV 已提交
105
    parser.add_argument("--cpu_threads", type=int, default=10)
W
WenmuZhou 已提交
106
    parser.add_argument("--use_pdserving", type=str2bool, default=False)
L
LDOUBLEV 已提交
107
    parser.add_argument("--warmup", type=str2bool, default=True)
W
WenmuZhou 已提交
108

L
LDOUBLEV 已提交
109
    # multi-process
littletomatodonkey's avatar
littletomatodonkey 已提交
110
    parser.add_argument("--use_mp", type=str2bool, default=False)
111 112
    parser.add_argument("--total_process_num", type=int, default=1)
    parser.add_argument("--process_id", type=int, default=0)
W
WenmuZhou 已提交
113

littletomatodonkey's avatar
littletomatodonkey 已提交
114
    parser.add_argument("--benchmark", type=str2bool, default=False)
L
LDOUBLEV 已提交
115
    parser.add_argument("--save_log_path", type=str, default="./log_output/")
D
Double_V 已提交
116

W
WenmuZhou 已提交
117
    parser.add_argument("--show_log", type=str2bool, default=True)
W
WenmuZhou 已提交
118
    return parser
W
WenmuZhou 已提交
119

120

121
def parse_args():
W
WenmuZhou 已提交
122
    parser = init_args()
L
LDOUBLEV 已提交
123 124 125
    return parser.parse_args()


W
WenmuZhou 已提交
126 127 128 129 130
def create_predictor(args, mode, logger):
    if mode == "det":
        model_dir = args.det_model_dir
    elif mode == 'cls':
        model_dir = args.cls_model_dir
J
Jethong 已提交
131
    elif mode == 'rec':
W
WenmuZhou 已提交
132
        model_dir = args.rec_model_dir
W
WenmuZhou 已提交
133 134
    elif mode == 'table':
        model_dir = args.table_model_dir
J
Jethong 已提交
135 136
    else:
        model_dir = args.e2e_model_dir
W
WenmuZhou 已提交
137 138 139 140

    if model_dir is None:
        logger.info("not find {} model file path {}".format(mode, model_dir))
        sys.exit(0)
文幕地方's avatar
文幕地方 已提交
141 142
    model_file_path = model_dir + "/inference.pdmodel"
    params_file_path = model_dir + "/inference.pdiparams"
W
WenmuZhou 已提交
143
    if not os.path.exists(model_file_path):
L
LDOUBLEV 已提交
144
        raise ValueError("not find model file path {}".format(model_file_path))
W
WenmuZhou 已提交
145
    if not os.path.exists(params_file_path):
L
LDOUBLEV 已提交
146 147
        raise ValueError("not find params file path {}".format(
            params_file_path))
W
WenmuZhou 已提交
148

W
WenmuZhou 已提交
149
    config = inference.Config(model_file_path, params_file_path)
W
WenmuZhou 已提交
150

L
LDOUBLEV 已提交
151 152 153 154 155 156 157 158 159 160
    if hasattr(args, 'precision'):
        if args.precision == "fp16" and args.use_tensorrt:
            precision = inference.PrecisionType.Half
        elif args.precision == "int8":
            precision = inference.PrecisionType.Int8
        else:
            precision = inference.PrecisionType.Float32
    else:
        precision = inference.PrecisionType.Float32

W
WenmuZhou 已提交
161
    if args.use_gpu:
162 163 164 165 166
        gpu_id = get_infer_gpuid()
        if gpu_id is None:
            raise ValueError(
                "Not found GPU in current device. Please check your device or set args.use_gpu as False"
            )
W
WenmuZhou 已提交
167
        config.enable_use_gpu(args.gpu_mem, 0)
L
LDOUBLEV 已提交
168 169
        if args.use_tensorrt:
            config.enable_tensorrt_engine(
D
Double_V 已提交
170
                precision_mode=precision,
L
LDOUBLEV 已提交
171
                max_batch_size=args.max_batch_size,
L
LDOUBLEV 已提交
172 173
                min_subgraph_size=args.min_subgraph_size)
            # skip the minmum trt subgraph
L
LDOUBLEV 已提交
174
        if mode == "det":
L
LDOUBLEV 已提交
175 176
            min_input_shape = {
                "x": [1, 3, 50, 50],
F
fengshuai03 已提交
177 178
                "conv2d_92.tmp_0": [1, 120, 20, 20],
                "conv2d_91.tmp_0": [1, 24, 10, 10],
L
LDOUBLEV 已提交
179
                "conv2d_59.tmp_0": [1, 96, 20, 20],
F
fengshuai03 已提交
180 181 182 183 184 185
                "nearest_interp_v2_1.tmp_0": [1, 256, 10, 10],
                "nearest_interp_v2_2.tmp_0": [1, 256, 20, 20],
                "conv2d_124.tmp_0": [1, 256, 20, 20],
                "nearest_interp_v2_3.tmp_0": [1, 64, 20, 20],
                "nearest_interp_v2_4.tmp_0": [1, 64, 20, 20],
                "nearest_interp_v2_5.tmp_0": [1, 64, 20, 20],
L
LDOUBLEV 已提交
186
                "elementwise_add_7": [1, 56, 2, 2],
F
fengshuai03 已提交
187
                "nearest_interp_v2_0.tmp_0": [1, 256, 2, 2]
L
LDOUBLEV 已提交
188 189 190
            }
            max_input_shape = {
                "x": [1, 3, 2000, 2000],
F
fengshuai03 已提交
191 192
                "conv2d_92.tmp_0": [1, 120, 400, 400],
                "conv2d_91.tmp_0": [1, 24, 200, 200],
L
LDOUBLEV 已提交
193
                "conv2d_59.tmp_0": [1, 96, 400, 400],
F
fengshuai03 已提交
194
                "nearest_interp_v2_1.tmp_0": [1, 256, 200, 200],
L
LDOUBLEV 已提交
195
                "conv2d_124.tmp_0": [1, 256, 400, 400],
F
fengshuai03 已提交
196 197 198 199
                "nearest_interp_v2_2.tmp_0": [1, 256, 400, 400],
                "nearest_interp_v2_3.tmp_0": [1, 64, 400, 400],
                "nearest_interp_v2_4.tmp_0": [1, 64, 400, 400],
                "nearest_interp_v2_5.tmp_0": [1, 64, 400, 400],
L
LDOUBLEV 已提交
200
                "elementwise_add_7": [1, 56, 400, 400],
F
fengshuai03 已提交
201
                "nearest_interp_v2_0.tmp_0": [1, 256, 400, 400]
L
LDOUBLEV 已提交
202 203 204
            }
            opt_input_shape = {
                "x": [1, 3, 640, 640],
F
fengshuai03 已提交
205 206
                "conv2d_92.tmp_0": [1, 120, 160, 160],
                "conv2d_91.tmp_0": [1, 24, 80, 80],
L
LDOUBLEV 已提交
207
                "conv2d_59.tmp_0": [1, 96, 160, 160],
F
fengshuai03 已提交
208 209
                "nearest_interp_v2_1.tmp_0": [1, 256, 80, 80],
                "nearest_interp_v2_2.tmp_0": [1, 256, 160, 160],
L
LDOUBLEV 已提交
210
                "conv2d_124.tmp_0": [1, 256, 160, 160],
F
fengshuai03 已提交
211 212 213
                "nearest_interp_v2_3.tmp_0": [1, 64, 160, 160],
                "nearest_interp_v2_4.tmp_0": [1, 64, 160, 160],
                "nearest_interp_v2_5.tmp_0": [1, 64, 160, 160],
L
LDOUBLEV 已提交
214
                "elementwise_add_7": [1, 56, 40, 40],
F
fengshuai03 已提交
215
                "nearest_interp_v2_0.tmp_0": [1, 256, 40, 40]
L
LDOUBLEV 已提交
216
            }
F
fengshuai03 已提交
217
            min_pact_shape = {
littletomatodonkey's avatar
littletomatodonkey 已提交
218 219 220 221
                "nearest_interp_v2_26.tmp_0": [1, 256, 20, 20],
                "nearest_interp_v2_27.tmp_0": [1, 64, 20, 20],
                "nearest_interp_v2_28.tmp_0": [1, 64, 20, 20],
                "nearest_interp_v2_29.tmp_0": [1, 64, 20, 20]
F
fengshuai03 已提交
222 223
            }
            max_pact_shape = {
littletomatodonkey's avatar
littletomatodonkey 已提交
224 225 226 227
                "nearest_interp_v2_26.tmp_0": [1, 256, 400, 400],
                "nearest_interp_v2_27.tmp_0": [1, 64, 400, 400],
                "nearest_interp_v2_28.tmp_0": [1, 64, 400, 400],
                "nearest_interp_v2_29.tmp_0": [1, 64, 400, 400]
F
fengshuai03 已提交
228 229
            }
            opt_pact_shape = {
littletomatodonkey's avatar
littletomatodonkey 已提交
230 231 232 233
                "nearest_interp_v2_26.tmp_0": [1, 256, 160, 160],
                "nearest_interp_v2_27.tmp_0": [1, 64, 160, 160],
                "nearest_interp_v2_28.tmp_0": [1, 64, 160, 160],
                "nearest_interp_v2_29.tmp_0": [1, 64, 160, 160]
F
fengshuai03 已提交
234 235 236 237
            }
            min_input_shape.update(min_pact_shape)
            max_input_shape.update(max_pact_shape)
            opt_input_shape.update(opt_pact_shape)
L
LDOUBLEV 已提交
238 239 240 241 242 243 244 245
        elif mode == "rec":
            min_input_shape = {"x": [args.rec_batch_num, 3, 32, 10]}
            max_input_shape = {"x": [args.rec_batch_num, 3, 32, 2000]}
            opt_input_shape = {"x": [args.rec_batch_num, 3, 32, 320]}
        elif mode == "cls":
            min_input_shape = {"x": [args.rec_batch_num, 3, 48, 10]}
            max_input_shape = {"x": [args.rec_batch_num, 3, 48, 2000]}
            opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]}
L
LDOUBLEV 已提交
246 247 248 249
        else:
            min_input_shape = {"x": [1, 3, 10, 10]}
            max_input_shape = {"x": [1, 3, 1000, 1000]}
            opt_input_shape = {"x": [1, 3, 500, 500]}
L
LDOUBLEV 已提交
250 251 252
        config.set_trt_dynamic_shape_info(min_input_shape, max_input_shape,
                                          opt_input_shape)

W
WenmuZhou 已提交
253 254
    else:
        config.disable_gpu()
L
LDOUBLEV 已提交
255 256 257
        if hasattr(args, "cpu_threads"):
            config.set_cpu_math_library_num_threads(args.cpu_threads)
        else:
W
WenmuZhou 已提交
258
            # default cpu threads as 10
L
LDOUBLEV 已提交
259
            config.set_cpu_math_library_num_threads(10)
W
WenmuZhou 已提交
260 261 262 263 264
        if args.enable_mkldnn:
            # cache 10 different shapes for mkldnn to avoid memory leak
            config.set_mkldnn_cache_capacity(10)
            config.enable_mkldnn()

L
LDOUBLEV 已提交
265 266
    # enable memory optim
    config.enable_memory_optim()
L
LDOUBLEV 已提交
267
    #config.disable_glog_info()
W
WenmuZhou 已提交
268

W
WenmuZhou 已提交
269
    config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
W
WenmuZhou 已提交
270
    if mode == 'table':
W
WenmuZhou 已提交
271
        config.delete_pass("fc_fuse_pass")  # not supported for table
W
WenmuZhou 已提交
272
    config.switch_use_feed_fetch_ops(False)
W
WenmuZhou 已提交
273
    config.switch_ir_optim(True)
274

W
WenmuZhou 已提交
275 276
    # create predictor
    predictor = inference.create_predictor(config)
W
WenmuZhou 已提交
277 278
    input_names = predictor.get_input_names()
    for name in input_names:
W
WenmuZhou 已提交
279
        input_tensor = predictor.get_input_handle(name)
W
WenmuZhou 已提交
280 281 282
    output_names = predictor.get_output_names()
    output_tensors = []
    for output_name in output_names:
W
WenmuZhou 已提交
283
        output_tensor = predictor.get_output_handle(output_name)
W
WenmuZhou 已提交
284
        output_tensors.append(output_tensor)
L
LDOUBLEV 已提交
285
    return predictor, input_tensor, output_tensors, config
W
WenmuZhou 已提交
286 287


L
LDOUBLEV 已提交
288 289 290 291 292 293 294 295 296 297 298 299 300 301
def get_infer_gpuid():
    cmd = "nvidia-smi"
    res = os.popen(cmd).readlines()
    if len(res) == 0:
        return None
    cmd = "env | grep CUDA_VISIBLE_DEVICES"
    env_cuda = os.popen(cmd).readlines()
    if len(env_cuda) == 0:
        return 0
    else:
        gpu_id = env_cuda[0].strip().split("=")[1]
        return int(gpu_id[0])


J
Jethong 已提交
302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317
def draw_e2e_res(dt_boxes, strs, img_path):
    src_im = cv2.imread(img_path)
    for box, str in zip(dt_boxes, strs):
        box = box.astype(np.int32).reshape((-1, 1, 2))
        cv2.polylines(src_im, [box], True, color=(255, 255, 0), thickness=2)
        cv2.putText(
            src_im,
            str,
            org=(int(box[0, 0, 0]), int(box[0, 0, 1])),
            fontFace=cv2.FONT_HERSHEY_COMPLEX,
            fontScale=0.7,
            color=(0, 255, 0),
            thickness=1)
    return src_im


L
LDOUBLEV 已提交
318
def draw_text_det_res(dt_boxes, img_path):
L
LDOUBLEV 已提交
319 320 321 322
    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)
L
LDOUBLEV 已提交
323
    return src_im
L
LDOUBLEV 已提交
324 325


L
LDOUBLEV 已提交
326 327
def resize_img(img, input_size=600):
    """
L
LDOUBLEV 已提交
328
    resize img and limit the longest side of the image to input_size
L
LDOUBLEV 已提交
329 330 331 332 333
    """
    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)
W
WenmuZhou 已提交
334 335
    img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
    return img
L
LDOUBLEV 已提交
336 337


W
WenmuZhou 已提交
338 339 340 341 342
def draw_ocr(image,
             boxes,
             txts=None,
             scores=None,
             drop_score=0.5,
L
LDOUBLEV 已提交
343
             font_path="./doc/fonts/simfang.ttf"):
344 345 346
    """
    Visualize the results of OCR detection and recognition
    args:
L
LDOUBLEV 已提交
347
        image(Image|array): RGB image
348 349 350 351
        boxes(list): boxes with shape(N, 4, 2)
        txts(list): the texts
        scores(list): txxs corresponding scores
        drop_score(float): only scores greater than drop_threshold will be visualized
W
WenmuZhou 已提交
352
        font_path: the path of font which is used to draw text
353 354 355
    return(array):
        the visualized img
    """
L
LDOUBLEV 已提交
356 357
    if scores is None:
        scores = [1] * len(boxes)
W
WenmuZhou 已提交
358 359 360 361
    box_num = len(boxes)
    for i in range(box_num):
        if scores is not None and (scores[i] < drop_score or
                                   math.isnan(scores[i])):
L
LDOUBLEV 已提交
362
            continue
W
WenmuZhou 已提交
363
        box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64)
L
LDOUBLEV 已提交
364
        image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
W
WenmuZhou 已提交
365
    if txts is not None:
L
LDOUBLEV 已提交
366
        img = np.array(resize_img(image, input_size=600))
367
        txt_img = text_visual(
W
WenmuZhou 已提交
368 369 370 371 372 373
            txts,
            scores,
            img_h=img.shape[0],
            img_w=600,
            threshold=drop_score,
            font_path=font_path)
374
        img = np.concatenate([np.array(img), np.array(txt_img)], axis=1)
L
LDOUBLEV 已提交
375 376
        return img
    return image
377 378


W
WenmuZhou 已提交
379 380 381 382 383 384
def draw_ocr_box_txt(image,
                     boxes,
                     txts,
                     scores=None,
                     drop_score=0.5,
                     font_path="./doc/simfang.ttf"):
385 386 387
    h, w = image.height, image.width
    img_left = image.copy()
    img_right = Image.new('RGB', (w, h), (255, 255, 255))
388 389

    import random
L
LDOUBLEV 已提交
390

391 392 393
    random.seed(0)
    draw_left = ImageDraw.Draw(img_left)
    draw_right = ImageDraw.Draw(img_right)
W
WenmuZhou 已提交
394 395 396
    for idx, (box, txt) in enumerate(zip(boxes, txts)):
        if scores is not None and scores[idx] < drop_score:
            continue
T
tink2123 已提交
397 398
        color = (random.randint(0, 255), random.randint(0, 255),
                 random.randint(0, 255))
399
        draw_left.polygon(box, fill=color)
T
tink2123 已提交
400 401 402 403 404 405 406 407 408 409
        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)
410 411
        if box_height > 2 * box_width:
            font_size = max(int(box_width * 0.9), 10)
W
WenmuZhou 已提交
412
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
413 414 415
            cur_y = box[0][1]
            for c in txt:
                char_size = font.getsize(c)
T
tink2123 已提交
416 417
                draw_right.text(
                    (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
418 419 420
                cur_y += char_size[1]
        else:
            font_size = max(int(box_height * 0.8), 10)
W
WenmuZhou 已提交
421
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
T
tink2123 已提交
422 423
            draw_right.text(
                [box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
424 425 426 427
    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))
428 429 430
    return np.array(img_show)


431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454
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)


W
WenmuZhou 已提交
455 456 457 458 459 460
def text_visual(texts,
                scores,
                img_h=400,
                img_w=600,
                threshold=0.,
                font_path="./doc/simfang.ttf"):
461 462 463 464 465 466 467
    """
    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
W
WenmuZhou 已提交
468
        font_path: the path of font which is used to draw text
469 470 471 472 473 474 475 476 477
    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
L
LDOUBLEV 已提交
478 479
        blank_img = Image.fromarray(blank_img).convert("RGB")
        draw_txt = ImageDraw.Draw(blank_img)
480
        return blank_img, draw_txt
L
LDOUBLEV 已提交
481

482 483 484 485
    blank_img, draw_txt = create_blank_img()

    font_size = 20
    txt_color = (0, 0, 0)
W
WenmuZhou 已提交
486
    font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
487 488 489

    gap = font_size + 5
    txt_img_list = []
L
LDOUBLEV 已提交
490
    count, index = 1, 0
491 492
    for idx, txt in enumerate(texts):
        index += 1
L
LDOUBLEV 已提交
493
        if scores[idx] < threshold or math.isnan(scores[idx]):
494 495 496 497 498 499 500 501 502 503 504
            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
L
LDOUBLEV 已提交
505
            draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
506 507 508 509 510
            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
L
LDOUBLEV 已提交
511
            count += 1
512 513 514
        if first_line:
            new_txt = str(index) + ': ' + txt + '   ' + '%.3f' % (scores[idx])
        else:
L
LDOUBLEV 已提交
515
            new_txt = "  " + txt + "  " + '%.3f' % (scores[idx])
L
LDOUBLEV 已提交
516
        draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
517
        # whether add new blank img or not
L
LDOUBLEV 已提交
518
        if count >= img_h // gap - 1 and idx + 1 < len(texts):
519 520 521
            txt_img_list.append(np.array(blank_img))
            blank_img, draw_txt = create_blank_img()
            count = 0
L
LDOUBLEV 已提交
522
        count += 1
523 524 525 526 527 528
    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)
L
LDOUBLEV 已提交
529 530


D
dyning 已提交
531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549
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


W
WenmuZhou 已提交
550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584
def get_rotate_crop_image(img, points):
    '''
    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
    '''
    assert len(points) == 4, "shape of points must be 4*2"
    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],
                          [img_crop_width, img_crop_height],
                          [0, img_crop_height]])
    M = cv2.getPerspectiveTransform(points, pts_std)
    dst_img = cv2.warpPerspective(
        img,
        M, (img_crop_width, img_crop_height),
        borderMode=cv2.BORDER_REPLICATE,
        flags=cv2.INTER_CUBIC)
    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


L
LDOUBLEV 已提交
585
if __name__ == '__main__':
L
LDOUBLEV 已提交
586
    pass