utility.py 22.6 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
文幕地方's avatar
文幕地方 已提交
18
import platform
L
LDOUBLEV 已提交
19 20
import cv2
import numpy as np
Z
zhoujun 已提交
21
import paddle
L
LDOUBLEV 已提交
22
from PIL import Image, ImageDraw, ImageFont
23
import math
W
WenmuZhou 已提交
24
from paddle import inference
L
LDOUBLEV 已提交
25 26
import time
from ppocr.utils.logging import get_logger
W
WenmuZhou 已提交
27

L
LDOUBLEV 已提交
28

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


W
WenmuZhou 已提交
33
def init_args():
L
LDOUBLEV 已提交
34
    parser = argparse.ArgumentParser()
W
WenmuZhou 已提交
35
    # params for prediction engine
L
LDOUBLEV 已提交
36
    parser.add_argument("--use_gpu", type=str2bool, default=True)
X
xiaoting 已提交
37
    parser.add_argument("--use_xpu", type=str2bool, default=False)
L
LDOUBLEV 已提交
38 39
    parser.add_argument("--ir_optim", type=str2bool, default=True)
    parser.add_argument("--use_tensorrt", type=str2bool, default=False)
L
LDOUBLEV 已提交
40
    parser.add_argument("--min_subgraph_size", type=int, default=15)
L
LDOUBLEV 已提交
41
    parser.add_argument("--precision", type=str, default="fp32")
L
LDOUBLEV 已提交
42
    parser.add_argument("--gpu_mem", type=int, default=500)
L
LDOUBLEV 已提交
43

W
WenmuZhou 已提交
44
    # params for text detector
L
LDOUBLEV 已提交
45 46 47
    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 已提交
48 49
    parser.add_argument("--det_limit_side_len", type=float, default=960)
    parser.add_argument("--det_limit_type", type=str, default='max')
L
LDOUBLEV 已提交
50

W
WenmuZhou 已提交
51
    # DB parmas
L
LDOUBLEV 已提交
52
    parser.add_argument("--det_db_thresh", type=float, default=0.3)
L
LDOUBLEV 已提交
53 54
    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 已提交
55
    parser.add_argument("--max_batch_size", type=int, default=10)
littletomatodonkey's avatar
littletomatodonkey 已提交
56
    parser.add_argument("--use_dilation", type=str2bool, default=False)
littletomatodonkey's avatar
littletomatodonkey 已提交
57
    parser.add_argument("--det_db_score_mode", type=str, default="fast")
W
WenmuZhou 已提交
58
    # EAST parmas
L
LDOUBLEV 已提交
59 60 61 62
    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 已提交
63
    # SAST parmas
L
licx 已提交
64 65
    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 已提交
66
    parser.add_argument("--det_sast_polygon", type=str2bool, default=False)
L
licx 已提交
67

W
WenmuZhou 已提交
68 69 70 71
    # 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)
文幕地方's avatar
文幕地方 已提交
72
    parser.add_argument("--det_pse_box_type", type=str, default='quad')
W
WenmuZhou 已提交
73 74
    parser.add_argument("--det_pse_scale", type=int, default=1)

文幕地方's avatar
文幕地方 已提交
75 76 77 78 79 80 81
    # FCE parmas
    parser.add_argument("--scales", type=list, default=[8, 16, 32])
    parser.add_argument("--alpha", type=float, default=1.0)
    parser.add_argument("--beta", type=float, default=1.0)
    parser.add_argument("--fourier_degree", type=int, default=5)
    parser.add_argument("--det_fce_box_type", type=str, default='poly')

W
WenmuZhou 已提交
82
    # params for text recognizer
A
andyjpaddle 已提交
83
    parser.add_argument("--rec_algorithm", type=str, default='SVTR_LCNet')
L
LDOUBLEV 已提交
84
    parser.add_argument("--rec_model_dir", type=str)
T
tink2123 已提交
85
    parser.add_argument("--rec_image_shape", type=str, default="3, 48, 320")
L
LDOUBLEV 已提交
86
    parser.add_argument("--rec_batch_num", type=int, default=6)
T
fix bug  
tink2123 已提交
87
    parser.add_argument("--max_text_length", type=int, default=25)
L
LDOUBLEV 已提交
88 89 90 91
    parser.add_argument(
        "--rec_char_dict_path",
        type=str,
        default="./ppocr/utils/ppocr_keys_v1.txt")
W
WenmuZhou 已提交
92 93
    parser.add_argument("--use_space_char", type=str2bool, default=True)
    parser.add_argument(
T
tink2123 已提交
94
        "--vis_font_path", type=str, default="./doc/fonts/simfang.ttf")
W
WenmuZhou 已提交
95
    parser.add_argument("--drop_score", type=float, default=0.5)
W
WenmuZhou 已提交
96

J
Jethong 已提交
97 98 99 100 101 102 103 104 105
    # 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 已提交
106
        "--e2e_char_dict_path", type=str, default="./ppocr/utils/ic15_dict.txt")
J
Jethong 已提交
107
    parser.add_argument("--e2e_pgnet_valid_set", type=str, default='totaltext')
J
Jethong 已提交
108
    parser.add_argument("--e2e_pgnet_mode", type=str, default='fast')
J
Jethong 已提交
109

W
WenmuZhou 已提交
110 111 112 113 114
    # 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 已提交
115
    parser.add_argument("--cls_batch_num", type=int, default=6)
W
WenmuZhou 已提交
116 117 118
    parser.add_argument("--cls_thresh", type=float, default=0.9)

    parser.add_argument("--enable_mkldnn", type=str2bool, default=False)
L
LDOUBLEV 已提交
119
    parser.add_argument("--cpu_threads", type=int, default=10)
W
WenmuZhou 已提交
120
    parser.add_argument("--use_pdserving", type=str2bool, default=False)
121 122
    parser.add_argument("--warmup", type=str2bool, default=False)

X
xiaoting 已提交
123 124 125 126 127
    # SR parmas
    parser.add_argument("--sr_model_dir", type=str)
    parser.add_argument("--sr_image_shape", type=str, default="3, 32, 128")
    parser.add_argument("--sr_batch_num", type=int, default=1)

128 129 130 131 132
    #
    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 已提交
133

L
LDOUBLEV 已提交
134
    # multi-process
littletomatodonkey's avatar
littletomatodonkey 已提交
135
    parser.add_argument("--use_mp", type=str2bool, default=False)
136 137
    parser.add_argument("--total_process_num", type=int, default=1)
    parser.add_argument("--process_id", type=int, default=0)
W
WenmuZhou 已提交
138

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

W
WenmuZhou 已提交
142
    parser.add_argument("--show_log", type=str2bool, default=True)
T
tink2123 已提交
143
    parser.add_argument("--use_onnx", type=str2bool, default=False)
W
WenmuZhou 已提交
144
    return parser
W
WenmuZhou 已提交
145

146

147
def parse_args():
W
WenmuZhou 已提交
148
    parser = init_args()
L
LDOUBLEV 已提交
149 150 151
    return parser.parse_args()


W
WenmuZhou 已提交
152 153 154 155 156
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 已提交
157
    elif mode == 'rec':
W
WenmuZhou 已提交
158
        model_dir = args.rec_model_dir
W
WenmuZhou 已提交
159 160
    elif mode == 'table':
        model_dir = args.table_model_dir
161 162
    elif mode == 'ser':
        model_dir = args.ser_model_dir
X
xiaoting 已提交
163 164
    elif mode == "sr":
        model_dir = args.sr_model_dir
文幕地方's avatar
文幕地方 已提交
165 166
    elif mode == 'layout':
        model_dir = args.layout_model_dir
J
Jethong 已提交
167 168
    else:
        model_dir = args.e2e_model_dir
W
WenmuZhou 已提交
169 170 171 172

    if model_dir is None:
        logger.info("not find {} model file path {}".format(mode, model_dir))
        sys.exit(0)
T
tink2123 已提交
173 174 175 176 177 178 179 180
    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 已提交
181

L
LDOUBLEV 已提交
182
    else:
Z
zhoujun 已提交
183 184 185 186 187 188 189
        file_names = ['model', 'inference']
        for file_name in file_names:
            model_file_path = '{}/{}.pdmodel'.format(model_dir, file_name)
            params_file_path = '{}/{}.pdiparams'.format(model_dir, file_name)
            if os.path.exists(model_file_path) and os.path.exists(
                    params_file_path):
                break
T
tink2123 已提交
190
        if not os.path.exists(model_file_path):
Z
zhoujun 已提交
191 192 193
            raise ValueError(
                "not find model.pdmodel or inference.pdmodel in {}".format(
                    model_dir))
T
tink2123 已提交
194
        if not os.path.exists(params_file_path):
Z
zhoujun 已提交
195 196 197
            raise ValueError(
                "not find model.pdiparams or inference.pdiparams in {}".format(
                    model_dir))
T
tink2123 已提交
198 199 200 201 202 203 204 205 206 207

        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 已提交
208
        else:
T
tink2123 已提交
209 210 211 212 213
            precision = inference.PrecisionType.Float32

        if args.use_gpu:
            gpu_id = get_infer_gpuid()
            if gpu_id is None:
L
LDOUBLEV 已提交
214
                logger.warning(
215
                    "GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jetson."
T
tink2123 已提交
216 217 218 219
                )
            config.enable_use_gpu(args.gpu_mem, 0)
            if args.use_tensorrt:
                config.enable_tensorrt_engine(
L
LDOUBLEV 已提交
220
                    workspace_size=1 << 30,
T
tink2123 已提交
221 222
                    precision_mode=precision,
                    max_batch_size=args.max_batch_size,
223 224
                    min_subgraph_size=args.
                    min_subgraph_size,  # skip the minmum trt subgraph
225
                    use_calib_mode=False)
226

littletomatodonkey's avatar
littletomatodonkey 已提交
227
                # collect shape
228 229
                trt_shape_f = os.path.join(model_dir,
                                           f"{mode}_trt_dynamic_shape.txt")
D
Double_V 已提交
230 231 232 233 234

                if not os.path.exists(trt_shape_f):
                    config.collect_shape_range_info(trt_shape_f)
                    logger.info(
                        f"collect dynamic shape info into : {trt_shape_f}")
littletomatodonkey's avatar
littletomatodonkey 已提交
235
                else:
236
                    logger.info(
D
Double_V 已提交
237
                        f"dynamic shape info file( {trt_shape_f} ) already exists, not need to generate again."
238
                    )
D
Double_V 已提交
239 240 241 242 243 244
                try:
                    config.enable_tuned_tensorrt_dynamic_shape(trt_shape_f,
                                                               True)
                except Exception as E:
                    logger.info(E)
                    logger.info("Please keep your paddlepaddle-gpu >= 2.3.0!")
L
LDOUBLEV 已提交
245

X
xiaoting 已提交
246 247
        elif args.use_xpu:
            config.enable_xpu(10 * 1024 * 1024)
L
LDOUBLEV 已提交
248
        else:
T
tink2123 已提交
249 250 251 252 253 254 255
            config.disable_gpu()
            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()
A
andyjpaddle 已提交
256 257 258 259 260
                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)
T
tink2123 已提交
261 262
        # enable memory optim
        config.enable_memory_optim()
littletomatodonkey's avatar
fix  
littletomatodonkey 已提交
263
        config.disable_glog_info()
T
tink2123 已提交
264
        config.delete_pass("conv_transpose_eltwiseadd_bn_fuse_pass")
T
tink2123 已提交
265
        config.delete_pass("matmul_transpose_reshape_fuse_pass")
T
tink2123 已提交
266 267 268 269 270 271 272 273
        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()
文幕地方's avatar
文幕地方 已提交
274
        if mode in ['ser', 're']:
275 276 277 278 279 280
            input_tensor = []
            for name in input_names:
                input_tensor.append(predictor.get_input_handle(name))
        else:
            for name in input_names:
                input_tensor = predictor.get_input_handle(name)
L
LDOUBLEV 已提交
281 282 283 284 285 286 287
        output_tensors = get_output_tensors(args, mode, predictor)
        return predictor, input_tensor, output_tensors, config


def get_output_tensors(args, mode, predictor):
    output_names = predictor.get_output_names()
    output_tensors = []
A
andyjpaddle 已提交
288
    if mode == "rec" and args.rec_algorithm in ["CRNN", "SVTR_LCNet"]:
L
LDOUBLEV 已提交
289 290 291
        output_name = 'softmax_0.tmp_0'
        if output_name in output_names:
            return [predictor.get_output_handle(output_name)]
L
LDOUBLEV 已提交
292 293 294 295
        else:
            for output_name in output_names:
                output_tensor = predictor.get_output_handle(output_name)
                output_tensors.append(output_tensor)
L
LDOUBLEV 已提交
296
    else:
T
tink2123 已提交
297 298 299
        for output_name in output_names:
            output_tensor = predictor.get_output_handle(output_name)
            output_tensors.append(output_tensor)
L
LDOUBLEV 已提交
300
    return output_tensors
W
WenmuZhou 已提交
301 302


L
LDOUBLEV 已提交
303
def get_infer_gpuid():
文幕地方's avatar
文幕地方 已提交
304 305 306 307
    sysstr = platform.system()
    if sysstr == "Windows":
        return 0

R
ronny1996 已提交
308 309 310 311
    if not paddle.fluid.core.is_compiled_with_rocm():
        cmd = "env | grep CUDA_VISIBLE_DEVICES"
    else:
        cmd = "env | grep HIP_VISIBLE_DEVICES"
L
LDOUBLEV 已提交
312 313 314 315 316 317 318 319
    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 已提交
320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335
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 已提交
336
def draw_text_det_res(dt_boxes, img_path):
L
LDOUBLEV 已提交
337 338 339 340
    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 已提交
341
    return src_im
L
LDOUBLEV 已提交
342 343


L
LDOUBLEV 已提交
344 345
def resize_img(img, input_size=600):
    """
L
LDOUBLEV 已提交
346
    resize img and limit the longest side of the image to input_size
L
LDOUBLEV 已提交
347 348 349 350 351
    """
    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 已提交
352 353
    img = cv2.resize(img, None, None, fx=im_scale, fy=im_scale)
    return img
L
LDOUBLEV 已提交
354 355


W
WenmuZhou 已提交
356 357 358 359 360
def draw_ocr(image,
             boxes,
             txts=None,
             scores=None,
             drop_score=0.5,
L
LDOUBLEV 已提交
361
             font_path="./doc/fonts/simfang.ttf"):
362 363 364
    """
    Visualize the results of OCR detection and recognition
    args:
L
LDOUBLEV 已提交
365
        image(Image|array): RGB image
366 367 368 369
        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 已提交
370
        font_path: the path of font which is used to draw text
371 372 373
    return(array):
        the visualized img
    """
L
LDOUBLEV 已提交
374 375
    if scores is None:
        scores = [1] * len(boxes)
W
WenmuZhou 已提交
376 377 378 379
    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 已提交
380
            continue
W
WenmuZhou 已提交
381
        box = np.reshape(np.array(boxes[i]), [-1, 1, 2]).astype(np.int64)
L
LDOUBLEV 已提交
382
        image = cv2.polylines(np.array(image), [box], True, (255, 0, 0), 2)
W
WenmuZhou 已提交
383
    if txts is not None:
L
LDOUBLEV 已提交
384
        img = np.array(resize_img(image, input_size=600))
385
        txt_img = text_visual(
W
WenmuZhou 已提交
386 387 388 389 390 391
            txts,
            scores,
            img_h=img.shape[0],
            img_w=600,
            threshold=drop_score,
            font_path=font_path)
392
        img = np.concatenate([np.array(img), np.array(txt_img)], axis=1)
L
LDOUBLEV 已提交
393 394
        return img
    return image
395 396


W
WenmuZhou 已提交
397 398 399 400 401 402
def draw_ocr_box_txt(image,
                     boxes,
                     txts,
                     scores=None,
                     drop_score=0.5,
                     font_path="./doc/simfang.ttf"):
403 404 405
    h, w = image.height, image.width
    img_left = image.copy()
    img_right = Image.new('RGB', (w, h), (255, 255, 255))
406 407

    import random
L
LDOUBLEV 已提交
408

409 410 411
    random.seed(0)
    draw_left = ImageDraw.Draw(img_left)
    draw_right = ImageDraw.Draw(img_right)
W
WenmuZhou 已提交
412 413 414
    for idx, (box, txt) in enumerate(zip(boxes, txts)):
        if scores is not None and scores[idx] < drop_score:
            continue
T
tink2123 已提交
415 416
        color = (random.randint(0, 255), random.randint(0, 255),
                 random.randint(0, 255))
417
        draw_left.polygon(box, fill=color)
T
tink2123 已提交
418 419 420 421 422 423 424 425 426 427
        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)
428 429
        if box_height > 2 * box_width:
            font_size = max(int(box_width * 0.9), 10)
W
WenmuZhou 已提交
430
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
431 432 433
            cur_y = box[0][1]
            for c in txt:
                char_size = font.getsize(c)
T
tink2123 已提交
434 435
                draw_right.text(
                    (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
436 437 438
                cur_y += char_size[1]
        else:
            font_size = max(int(box_height * 0.8), 10)
W
WenmuZhou 已提交
439
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
T
tink2123 已提交
440 441
            draw_right.text(
                [box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
442 443 444 445
    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))
446 447 448
    return np.array(img_show)


449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472
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 已提交
473 474 475 476 477 478
def text_visual(texts,
                scores,
                img_h=400,
                img_w=600,
                threshold=0.,
                font_path="./doc/simfang.ttf"):
479 480 481 482 483 484 485
    """
    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 已提交
486
        font_path: the path of font which is used to draw text
487 488 489 490 491 492 493 494 495
    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 已提交
496 497
        blank_img = Image.fromarray(blank_img).convert("RGB")
        draw_txt = ImageDraw.Draw(blank_img)
498
        return blank_img, draw_txt
L
LDOUBLEV 已提交
499

500 501 502 503
    blank_img, draw_txt = create_blank_img()

    font_size = 20
    txt_color = (0, 0, 0)
W
WenmuZhou 已提交
504
    font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
505 506 507

    gap = font_size + 5
    txt_img_list = []
L
LDOUBLEV 已提交
508
    count, index = 1, 0
509 510
    for idx, txt in enumerate(texts):
        index += 1
L
LDOUBLEV 已提交
511
        if scores[idx] < threshold or math.isnan(scores[idx]):
512 513 514 515 516 517 518 519 520 521 522
            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 已提交
523
            draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
524 525 526 527 528
            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 已提交
529
            count += 1
530 531 532
        if first_line:
            new_txt = str(index) + ': ' + txt + '   ' + '%.3f' % (scores[idx])
        else:
L
LDOUBLEV 已提交
533
            new_txt = "  " + txt + "  " + '%.3f' % (scores[idx])
L
LDOUBLEV 已提交
534
        draw_txt.text((0, gap * count), new_txt, txt_color, font=font)
535
        # whether add new blank img or not
L
LDOUBLEV 已提交
536
        if count >= img_h // gap - 1 and idx + 1 < len(texts):
537 538 539
            txt_img_list.append(np.array(blank_img))
            blank_img, draw_txt = create_blank_img()
            count = 0
L
LDOUBLEV 已提交
540
        count += 1
541 542 543 544 545 546
    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 已提交
547 548


D
dyning 已提交
549 550 551
def base64_to_cv2(b64str):
    import base64
    data = base64.b64decode(b64str.encode('utf8'))
552
    data = np.frombuffer(data, np.uint8)
D
dyning 已提交
553 554 555 556 557 558 559 560 561 562 563 564 565 566 567
    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 已提交
568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602
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 已提交
603 604 605 606 607 608
def check_gpu(use_gpu):
    if use_gpu and not paddle.is_compiled_with_cuda():
        use_gpu = False
    return use_gpu


L
LDOUBLEV 已提交
609
if __name__ == '__main__':
L
LDOUBLEV 已提交
610
    pass