utility.py 28.1 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
Z
zhoujun 已提交
20
import paddle
L
LDOUBLEV 已提交
21
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 67 68 69
    # PSE parmas
    parser.add_argument("--det_pse_thresh", type=float, default=0)
    parser.add_argument("--det_pse_box_thresh", type=float, default=0.85)
    parser.add_argument("--det_pse_min_area", type=float, default=16)
W
WenmuZhou 已提交
70
    parser.add_argument("--det_pse_box_type", type=str, default='box')
W
WenmuZhou 已提交
71 72
    parser.add_argument("--det_pse_scale", type=int, default=1)

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

J
Jethong 已提交
88 89 90 91 92 93 94 95 96
    # 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 已提交
97
        "--e2e_char_dict_path", type=str, default="./ppocr/utils/ic15_dict.txt")
J
Jethong 已提交
98
    parser.add_argument("--e2e_pgnet_valid_set", type=str, default='totaltext')
J
Jethong 已提交
99
    parser.add_argument("--e2e_pgnet_mode", type=str, default='fast')
J
Jethong 已提交
100

W
WenmuZhou 已提交
101 102 103 104 105
    # 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 已提交
106
    parser.add_argument("--cls_batch_num", type=int, default=6)
W
WenmuZhou 已提交
107 108 109
    parser.add_argument("--cls_thresh", type=float, default=0.9)

    parser.add_argument("--enable_mkldnn", type=str2bool, default=False)
L
LDOUBLEV 已提交
110
    parser.add_argument("--cpu_threads", type=int, default=10)
W
WenmuZhou 已提交
111
    parser.add_argument("--use_pdserving", type=str2bool, default=False)
112 113 114 115 116 117 118
    parser.add_argument("--warmup", type=str2bool, default=False)

    #
    parser.add_argument(
        "--draw_img_save_dir", type=str, default="./inference_results")
    parser.add_argument("--save_crop_res", type=str2bool, default=False)
    parser.add_argument("--crop_res_save_dir", type=str, default="./output")
W
WenmuZhou 已提交
119

L
LDOUBLEV 已提交
120
    # multi-process
littletomatodonkey's avatar
littletomatodonkey 已提交
121
    parser.add_argument("--use_mp", type=str2bool, default=False)
122 123
    parser.add_argument("--total_process_num", type=int, default=1)
    parser.add_argument("--process_id", type=int, default=0)
W
WenmuZhou 已提交
124

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

W
WenmuZhou 已提交
128
    parser.add_argument("--show_log", type=str2bool, default=True)
T
tink2123 已提交
129
    parser.add_argument("--use_onnx", type=str2bool, default=False)
W
WenmuZhou 已提交
130
    return parser
W
WenmuZhou 已提交
131

132

133
def parse_args():
W
WenmuZhou 已提交
134
    parser = init_args()
L
LDOUBLEV 已提交
135 136 137
    return parser.parse_args()


W
WenmuZhou 已提交
138 139 140 141 142
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 已提交
143
    elif mode == 'rec':
W
WenmuZhou 已提交
144
        model_dir = args.rec_model_dir
W
WenmuZhou 已提交
145 146
    elif mode == 'table':
        model_dir = args.table_model_dir
J
Jethong 已提交
147 148
    else:
        model_dir = args.e2e_model_dir
W
WenmuZhou 已提交
149 150 151 152

    if model_dir is None:
        logger.info("not find {} model file path {}".format(mode, model_dir))
        sys.exit(0)
T
tink2123 已提交
153 154 155 156 157 158 159 160
    if args.use_onnx:
        import onnxruntime as ort
        model_file_path = model_dir
        if not os.path.exists(model_file_path):
            raise ValueError("not find model file path {}".format(
                model_file_path))
        sess = ort.InferenceSession(model_file_path)
        return sess, sess.get_inputs()[0], None, None
L
LDOUBLEV 已提交
161

W
WenmuZhou 已提交
162
    if args.use_gpu:
163 164 165 166 167
        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 已提交
168
        config.enable_use_gpu(args.gpu_mem, 0)
L
LDOUBLEV 已提交
169 170
        if args.use_tensorrt:
            config.enable_tensorrt_engine(
L
LDOUBLEV 已提交
171
                workspace_size=1 << 30,
D
Double_V 已提交
172
                precision_mode=precision,
L
LDOUBLEV 已提交
173
                max_batch_size=args.max_batch_size,
L
LDOUBLEV 已提交
174 175
                min_subgraph_size=args.min_subgraph_size)
            # skip the minmum trt subgraph
L
LDOUBLEV 已提交
176
        if mode == "det":
L
LDOUBLEV 已提交
177 178
            min_input_shape = {
                "x": [1, 3, 50, 50],
F
fengshuai03 已提交
179 180
                "conv2d_92.tmp_0": [1, 120, 20, 20],
                "conv2d_91.tmp_0": [1, 24, 10, 10],
L
LDOUBLEV 已提交
181
                "conv2d_59.tmp_0": [1, 96, 20, 20],
F
fengshuai03 已提交
182 183 184 185 186 187
                "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 已提交
188
                "elementwise_add_7": [1, 56, 2, 2],
F
fengshuai03 已提交
189
                "nearest_interp_v2_0.tmp_0": [1, 256, 2, 2]
L
LDOUBLEV 已提交
190 191 192
            }
            max_input_shape = {
                "x": [1, 3, 2000, 2000],
F
fengshuai03 已提交
193 194
                "conv2d_92.tmp_0": [1, 120, 400, 400],
                "conv2d_91.tmp_0": [1, 24, 200, 200],
L
LDOUBLEV 已提交
195
                "conv2d_59.tmp_0": [1, 96, 400, 400],
F
fengshuai03 已提交
196
                "nearest_interp_v2_1.tmp_0": [1, 256, 200, 200],
L
LDOUBLEV 已提交
197
                "conv2d_124.tmp_0": [1, 256, 400, 400],
F
fengshuai03 已提交
198 199 200 201
                "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 已提交
202
                "elementwise_add_7": [1, 56, 400, 400],
F
fengshuai03 已提交
203
                "nearest_interp_v2_0.tmp_0": [1, 256, 400, 400]
L
LDOUBLEV 已提交
204 205 206
            }
            opt_input_shape = {
                "x": [1, 3, 640, 640],
F
fengshuai03 已提交
207 208
                "conv2d_92.tmp_0": [1, 120, 160, 160],
                "conv2d_91.tmp_0": [1, 24, 80, 80],
L
LDOUBLEV 已提交
209
                "conv2d_59.tmp_0": [1, 96, 160, 160],
F
fengshuai03 已提交
210 211
                "nearest_interp_v2_1.tmp_0": [1, 256, 80, 80],
                "nearest_interp_v2_2.tmp_0": [1, 256, 160, 160],
L
LDOUBLEV 已提交
212
                "conv2d_124.tmp_0": [1, 256, 160, 160],
F
fengshuai03 已提交
213 214 215
                "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 已提交
216
                "elementwise_add_7": [1, 56, 40, 40],
F
fengshuai03 已提交
217
                "nearest_interp_v2_0.tmp_0": [1, 256, 40, 40]
L
LDOUBLEV 已提交
218
            }
F
fengshuai03 已提交
219
            min_pact_shape = {
littletomatodonkey's avatar
littletomatodonkey 已提交
220 221 222 223
                "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 已提交
224 225
            }
            max_pact_shape = {
littletomatodonkey's avatar
littletomatodonkey 已提交
226 227 228 229
                "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 已提交
230 231
            }
            opt_pact_shape = {
littletomatodonkey's avatar
littletomatodonkey 已提交
232 233 234 235
                "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 已提交
236 237 238 239
            }
            min_input_shape.update(min_pact_shape)
            max_input_shape.update(max_pact_shape)
            opt_input_shape.update(opt_pact_shape)
L
LDOUBLEV 已提交
240
        elif mode == "rec":
L
LDOUBLEV 已提交
241
            min_input_shape = {"x": [1, 3, 32, 10]}
L
LDOUBLEV 已提交
242 243 244
            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":
L
LDOUBLEV 已提交
245
            min_input_shape = {"x": [1, 3, 48, 10]}
L
LDOUBLEV 已提交
246 247
            max_input_shape = {"x": [args.rec_batch_num, 3, 48, 2000]}
            opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]}
L
LDOUBLEV 已提交
248

T
tink2123 已提交
249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266
        model_file_path = model_dir + "/inference.pdmodel"
        params_file_path = model_dir + "/inference.pdiparams"
        if not os.path.exists(model_file_path):
            raise ValueError("not find model file path {}".format(
                model_file_path))
        if not os.path.exists(params_file_path):
            raise ValueError("not find params file path {}".format(
                params_file_path))

        config = inference.Config(model_file_path, params_file_path)

        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
L
LDOUBLEV 已提交
267
        else:
T
tink2123 已提交
268 269 270 271 272
            precision = inference.PrecisionType.Float32

        if args.use_gpu:
            gpu_id = get_infer_gpuid()
            if gpu_id is None:
L
LDOUBLEV 已提交
273
                logger.warning(
L
LDOUBLEV 已提交
274
                    "GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jeston."
T
tink2123 已提交
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298
                )
            config.enable_use_gpu(args.gpu_mem, 0)
            if args.use_tensorrt:
                config.enable_tensorrt_engine(
                    precision_mode=precision,
                    max_batch_size=args.max_batch_size,
                    min_subgraph_size=args.min_subgraph_size)
                # skip the minmum trt subgraph
            if mode == "det":
                min_input_shape = {
                    "x": [1, 3, 50, 50],
                    "conv2d_92.tmp_0": [1, 120, 20, 20],
                    "conv2d_91.tmp_0": [1, 24, 10, 10],
                    "conv2d_59.tmp_0": [1, 96, 20, 20],
                    "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],
                    "elementwise_add_7": [1, 56, 2, 2],
                    "nearest_interp_v2_0.tmp_0": [1, 256, 2, 2]
                }
                max_input_shape = {
L
LDOUBLEV 已提交
299
                    "x": [1, 3, 1280, 1280],
T
tink2123 已提交
300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348
                    "conv2d_92.tmp_0": [1, 120, 400, 400],
                    "conv2d_91.tmp_0": [1, 24, 200, 200],
                    "conv2d_59.tmp_0": [1, 96, 400, 400],
                    "nearest_interp_v2_1.tmp_0": [1, 256, 200, 200],
                    "conv2d_124.tmp_0": [1, 256, 400, 400],
                    "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],
                    "elementwise_add_7": [1, 56, 400, 400],
                    "nearest_interp_v2_0.tmp_0": [1, 256, 400, 400]
                }
                opt_input_shape = {
                    "x": [1, 3, 640, 640],
                    "conv2d_92.tmp_0": [1, 120, 160, 160],
                    "conv2d_91.tmp_0": [1, 24, 80, 80],
                    "conv2d_59.tmp_0": [1, 96, 160, 160],
                    "nearest_interp_v2_1.tmp_0": [1, 256, 80, 80],
                    "nearest_interp_v2_2.tmp_0": [1, 256, 160, 160],
                    "conv2d_124.tmp_0": [1, 256, 160, 160],
                    "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],
                    "elementwise_add_7": [1, 56, 40, 40],
                    "nearest_interp_v2_0.tmp_0": [1, 256, 40, 40]
                }
                min_pact_shape = {
                    "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]
                }
                max_pact_shape = {
                    "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]
                }
                opt_pact_shape = {
                    "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]
                }
                min_input_shape.update(min_pact_shape)
                max_input_shape.update(max_pact_shape)
                opt_input_shape.update(opt_pact_shape)
            elif mode == "rec":
                min_input_shape = {"x": [1, 3, 32, 10]}
L
LDOUBLEV 已提交
349
                max_input_shape = {"x": [args.rec_batch_num, 3, 32, 1024]}
T
tink2123 已提交
350 351 352
                opt_input_shape = {"x": [args.rec_batch_num, 3, 32, 320]}
            elif mode == "cls":
                min_input_shape = {"x": [1, 3, 48, 10]}
L
LDOUBLEV 已提交
353
                max_input_shape = {"x": [args.rec_batch_num, 3, 48, 1024]}
T
tink2123 已提交
354 355 356
                opt_input_shape = {"x": [args.rec_batch_num, 3, 48, 320]}
            else:
                min_input_shape = {"x": [1, 3, 10, 10]}
L
LDOUBLEV 已提交
357 358
                max_input_shape = {"x": [1, 3, 512, 512]}
                opt_input_shape = {"x": [1, 3, 256, 256]}
T
tink2123 已提交
359 360
            config.set_trt_dynamic_shape_info(min_input_shape, max_input_shape,
                                              opt_input_shape)
L
LDOUBLEV 已提交
361

L
LDOUBLEV 已提交
362
        else:
T
tink2123 已提交
363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395
            config.disable_gpu()
            if hasattr(args, "cpu_threads"):
                config.set_cpu_math_library_num_threads(args.cpu_threads)
            else:
                # default cpu threads as 10
                config.set_cpu_math_library_num_threads(10)
            if args.enable_mkldnn:
                # cache 10 different shapes for mkldnn to avoid memory leak
                config.set_mkldnn_cache_capacity(10)
                config.enable_mkldnn()
                if args.precision == "fp16":
                    config.enable_mkldnn_bfloat16()
        # enable memory optim
        config.enable_memory_optim()
        config.disable_glog_info()

        config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
        if mode == 'table':
            config.delete_pass("fc_fuse_pass")  # not supported for table
        config.switch_use_feed_fetch_ops(False)
        config.switch_ir_optim(True)

        # create predictor
        predictor = inference.create_predictor(config)
        input_names = predictor.get_input_names()
        for name in input_names:
            input_tensor = predictor.get_input_handle(name)
        output_names = predictor.get_output_names()
        output_tensors = []
        for output_name in output_names:
            output_tensor = predictor.get_output_handle(output_name)
            output_tensors.append(output_tensor)
        return predictor, input_tensor, output_tensors, config
W
WenmuZhou 已提交
396 397


L
LDOUBLEV 已提交
398 399
def get_infer_gpuid():
    cmd = "nvidia-smi"
L
LDOUBLEV 已提交
400 401 402 403
    try:
        res = os.popen(cmd).readlines()
    except:
        res = None
L
LDOUBLEV 已提交
404 405 406 407 408 409 410 411 412 413 414
    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 已提交
415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430
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 已提交
431
def draw_text_det_res(dt_boxes, img_path):
L
LDOUBLEV 已提交
432 433 434 435
    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 已提交
436
    return src_im
L
LDOUBLEV 已提交
437 438


L
LDOUBLEV 已提交
439 440
def resize_img(img, input_size=600):
    """
L
LDOUBLEV 已提交
441
    resize img and limit the longest side of the image to input_size
L
LDOUBLEV 已提交
442 443 444 445 446
    """
    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 已提交
447 448
    img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
    return img
L
LDOUBLEV 已提交
449 450


W
WenmuZhou 已提交
451 452 453 454 455
def draw_ocr(image,
             boxes,
             txts=None,
             scores=None,
             drop_score=0.5,
L
LDOUBLEV 已提交
456
             font_path="./doc/fonts/simfang.ttf"):
457 458 459
    """
    Visualize the results of OCR detection and recognition
    args:
L
LDOUBLEV 已提交
460
        image(Image|array): RGB image
461 462 463 464
        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 已提交
465
        font_path: the path of font which is used to draw text
466 467 468
    return(array):
        the visualized img
    """
L
LDOUBLEV 已提交
469 470
    if scores is None:
        scores = [1] * len(boxes)
W
WenmuZhou 已提交
471 472 473 474
    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 已提交
475
            continue
W
WenmuZhou 已提交
476
        box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64)
L
LDOUBLEV 已提交
477
        image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
W
WenmuZhou 已提交
478
    if txts is not None:
L
LDOUBLEV 已提交
479
        img = np.array(resize_img(image, input_size=600))
480
        txt_img = text_visual(
W
WenmuZhou 已提交
481 482 483 484 485 486
            txts,
            scores,
            img_h=img.shape[0],
            img_w=600,
            threshold=drop_score,
            font_path=font_path)
487
        img = np.concatenate([np.array(img), np.array(txt_img)], axis=1)
L
LDOUBLEV 已提交
488 489
        return img
    return image
490 491


W
WenmuZhou 已提交
492 493 494 495 496 497
def draw_ocr_box_txt(image,
                     boxes,
                     txts,
                     scores=None,
                     drop_score=0.5,
                     font_path="./doc/simfang.ttf"):
498 499 500
    h, w = image.height, image.width
    img_left = image.copy()
    img_right = Image.new('RGB', (w, h), (255, 255, 255))
501 502

    import random
L
LDOUBLEV 已提交
503

504 505 506
    random.seed(0)
    draw_left = ImageDraw.Draw(img_left)
    draw_right = ImageDraw.Draw(img_right)
W
WenmuZhou 已提交
507 508 509
    for idx, (box, txt) in enumerate(zip(boxes, txts)):
        if scores is not None and scores[idx] < drop_score:
            continue
T
tink2123 已提交
510 511
        color = (random.randint(0, 255), random.randint(0, 255),
                 random.randint(0, 255))
512
        draw_left.polygon(box, fill=color)
T
tink2123 已提交
513 514 515 516 517 518 519 520 521 522
        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)
523 524
        if box_height > 2 * box_width:
            font_size = max(int(box_width * 0.9), 10)
W
WenmuZhou 已提交
525
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
526 527 528
            cur_y = box[0][1]
            for c in txt:
                char_size = font.getsize(c)
T
tink2123 已提交
529 530
                draw_right.text(
                    (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
531 532 533
                cur_y += char_size[1]
        else:
            font_size = max(int(box_height * 0.8), 10)
W
WenmuZhou 已提交
534
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
T
tink2123 已提交
535 536
            draw_right.text(
                [box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
537 538 539 540
    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))
541 542 543
    return np.array(img_show)


544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567
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 已提交
568 569 570 571 572 573
def text_visual(texts,
                scores,
                img_h=400,
                img_w=600,
                threshold=0.,
                font_path="./doc/simfang.ttf"):
574 575 576 577 578 579 580
    """
    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 已提交
581
        font_path: the path of font which is used to draw text
582 583 584 585 586 587 588 589 590
    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 已提交
591 592
        blank_img = Image.fromarray(blank_img).convert("RGB")
        draw_txt = ImageDraw.Draw(blank_img)
593
        return blank_img, draw_txt
L
LDOUBLEV 已提交
594

595 596 597 598
    blank_img, draw_txt = create_blank_img()

    font_size = 20
    txt_color = (0, 0, 0)
W
WenmuZhou 已提交
599
    font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
600 601 602

    gap = font_size + 5
    txt_img_list = []
L
LDOUBLEV 已提交
603
    count, index = 1, 0
604 605
    for idx, txt in enumerate(texts):
        index += 1
L
LDOUBLEV 已提交
606
        if scores[idx] < threshold or math.isnan(scores[idx]):
607 608 609 610 611 612 613 614 615 616 617
            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 已提交
618
            draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
619 620 621 622 623
            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 已提交
624
            count += 1
625 626 627
        if first_line:
            new_txt = str(index) + ': ' + txt + '   ' + '%.3f' % (scores[idx])
        else:
L
LDOUBLEV 已提交
628
            new_txt = "  " + txt + "  " + '%.3f' % (scores[idx])
L
LDOUBLEV 已提交
629
        draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
630
        # whether add new blank img or not
L
LDOUBLEV 已提交
631
        if count >= img_h // gap - 1 and idx + 1 < len(texts):
632 633 634
            txt_img_list.append(np.array(blank_img))
            blank_img, draw_txt = create_blank_img()
            count = 0
L
LDOUBLEV 已提交
635
        count += 1
636 637 638 639 640 641
    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 已提交
642 643


D
dyning 已提交
644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662
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 已提交
663 664 665 666 667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697
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


Z
zhoujun 已提交
698 699 700 701 702 703 704
def check_gpu(use_gpu):
    if use_gpu and not paddle.is_compiled_with_cuda():

        use_gpu = False
    return use_gpu


L
LDOUBLEV 已提交
705
if __name__ == '__main__':
L
LDOUBLEV 已提交
706
    pass