paddleclas.py 25.4 KB
Newer Older
G
gaotingquan 已提交
1
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
C
chenziheng 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15
#
# 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 os
16
from typing import Union, Generator
T
Tingquan Gao 已提交
17 18
import argparse
import shutil
T
Tingquan Gao 已提交
19
import textwrap
T
Tingquan Gao 已提交
20 21 22
import tarfile
import requests
from functools import partial
T
Tingquan Gao 已提交
23
from difflib import SequenceMatcher
C
chenziheng 已提交
24 25 26 27

import cv2
import numpy as np
from tqdm import tqdm
T
Tingquan Gao 已提交
28
from prettytable import PrettyTable
29
import paddle
T
Tingquan Gao 已提交
30

31 32
from .ppcls.arch import backbone
from .ppcls.utils import logger
33

34 35 36
from .deploy.python.predict_cls import ClsPredictor
from .deploy.utils.get_image_list import get_image_list
from .deploy.utils import config
T
Tingquan Gao 已提交
37

38 39 40
# for the PaddleClas Project
from . import deploy
from . import ppcls
T
Tingquan Gao 已提交
41

42
# for building model with loading pretrained weights from backbone
G
gaotingquan 已提交
43
logger.init_logger()
44

T
Tingquan Gao 已提交
45
__all__ = ["PaddleClas"]
T
Tingquan Gao 已提交
46

C
chenziheng 已提交
47
BASE_DIR = os.path.expanduser("~/.paddleclas/")
T
Tingquan Gao 已提交
48 49
BASE_INFERENCE_MODEL_DIR = os.path.join(BASE_DIR, "inference_model")
BASE_IMAGES_DIR = os.path.join(BASE_DIR, "images")
50 51
IMN_MODEL_BASE_DOWNLOAD_URL = "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/inference/{}_infer.tar"
IMN_MODEL_SERIES = {
T
Tingquan Gao 已提交
52
    "AlexNet": ["AlexNet"],
G
gaotingquan 已提交
53 54 55 56 57
    "CSWinTransformer": [
        "CSWinTransformer_tiny_224", "CSWinTransformer_small_224",
        "CSWinTransformer_base_224", "CSWinTransformer_base_384",
        "CSWinTransformer_large_224", "CSWinTransformer_large_384"
    ],
T
Tingquan Gao 已提交
58 59
    "DarkNet": ["DarkNet53"],
    "DeiT": [
T
Tingquan Gao 已提交
60 61 62 63
        "DeiT_base_distilled_patch16_224", "DeiT_base_distilled_patch16_384",
        "DeiT_base_patch16_224", "DeiT_base_patch16_384",
        "DeiT_small_distilled_patch16_224", "DeiT_small_patch16_224",
        "DeiT_tiny_distilled_patch16_224", "DeiT_tiny_patch16_224"
T
Tingquan Gao 已提交
64 65 66 67 68
    ],
    "DenseNet": [
        "DenseNet121", "DenseNet161", "DenseNet169", "DenseNet201",
        "DenseNet264"
    ],
69 70 71 72
    "DLA": [
        "DLA46_c", "DLA60x_c", "DLA34", "DLA60", "DLA60x", "DLA102", "DLA102x",
        "DLA102x2", "DLA169"
    ],
T
Tingquan Gao 已提交
73 74 75 76 77 78
    "DPN": ["DPN68", "DPN92", "DPN98", "DPN107", "DPN131"],
    "EfficientNet": [
        "EfficientNetB0", "EfficientNetB0_small", "EfficientNetB1",
        "EfficientNetB2", "EfficientNetB3", "EfficientNetB4", "EfficientNetB5",
        "EfficientNetB6", "EfficientNetB7"
    ],
G
gaotingquan 已提交
79
    "ESNet": ["ESNet_x0_25", "ESNet_x0_5", "ESNet_x0_75", "ESNet_x1_0"],
T
Tingquan Gao 已提交
80 81
    "GhostNet":
    ["GhostNet_x0_5", "GhostNet_x1_0", "GhostNet_x1_3", "GhostNet_x1_3_ssld"],
82
    "HarDNet": ["HarDNet39_ds", "HarDNet68_ds", "HarDNet68", "HarDNet85"],
T
Tingquan Gao 已提交
83 84 85 86 87 88
    "HRNet": [
        "HRNet_W18_C", "HRNet_W30_C", "HRNet_W32_C", "HRNet_W40_C",
        "HRNet_W44_C", "HRNet_W48_C", "HRNet_W64_C", "HRNet_W18_C_ssld",
        "HRNet_W48_C_ssld"
    ],
    "Inception": ["GoogLeNet", "InceptionV3", "InceptionV4"],
G
gaotingquan 已提交
89 90
    "LeViT":
    ["LeViT_128S", "LeViT_128", "LeViT_192", "LeViT_256", "LeViT_384"],
G
gaotingquan 已提交
91
    "MixNet": ["MixNet_S", "MixNet_M", "MixNet_L"],
T
Tingquan Gao 已提交
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108
    "MobileNetV1": [
        "MobileNetV1_x0_25", "MobileNetV1_x0_5", "MobileNetV1_x0_75",
        "MobileNetV1", "MobileNetV1_ssld"
    ],
    "MobileNetV2": [
        "MobileNetV2_x0_25", "MobileNetV2_x0_5", "MobileNetV2_x0_75",
        "MobileNetV2", "MobileNetV2_x1_5", "MobileNetV2_x2_0",
        "MobileNetV2_ssld"
    ],
    "MobileNetV3": [
        "MobileNetV3_small_x0_35", "MobileNetV3_small_x0_5",
        "MobileNetV3_small_x0_75", "MobileNetV3_small_x1_0",
        "MobileNetV3_small_x1_25", "MobileNetV3_large_x0_35",
        "MobileNetV3_large_x0_5", "MobileNetV3_large_x0_75",
        "MobileNetV3_large_x1_0", "MobileNetV3_large_x1_25",
        "MobileNetV3_small_x1_0_ssld", "MobileNetV3_large_x1_0_ssld"
    ],
G
gaotingquan 已提交
109
    "MobileViT": ["MobileViT_XXS", "MobileViT_XS", "MobileViT_S"],
G
gaotingquan 已提交
110 111 112 113 114 115
    "PPHGNet": [
        "PPHGNet_tiny",
        "PPHGNet_small",
        "PPHGNet_tiny_ssld",
        "PPHGNet_small_ssld",
    ],
G
gaotingquan 已提交
116 117 118 119
    "PPLCNet": [
        "PPLCNet_x0_25", "PPLCNet_x0_35", "PPLCNet_x0_5", "PPLCNet_x0_75",
        "PPLCNet_x1_0", "PPLCNet_x1_5", "PPLCNet_x2_0", "PPLCNet_x2_5"
    ],
G
gaotingquan 已提交
120
    "PPLCNetV2": ["PPLCNetV2_base"],
G
gaotingquan 已提交
121 122 123 124
    "PVTV2": [
        "PVT_V2_B0", "PVT_V2_B1", "PVT_V2_B2", "PVT_V2_B2_Linear", "PVT_V2_B3",
        "PVT_V2_B4", "PVT_V2_B5"
    ],
125
    "RedNet": ["RedNet26", "RedNet38", "RedNet50", "RedNet101", "RedNet152"],
T
Tingquan Gao 已提交
126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149
    "RegNet": ["RegNetX_4GF"],
    "Res2Net": [
        "Res2Net50_14w_8s", "Res2Net50_26w_4s", "Res2Net50_vd_26w_4s",
        "Res2Net200_vd_26w_4s", "Res2Net101_vd_26w_4s",
        "Res2Net50_vd_26w_4s_ssld", "Res2Net101_vd_26w_4s_ssld",
        "Res2Net200_vd_26w_4s_ssld"
    ],
    "ResNeSt": ["ResNeSt50", "ResNeSt50_fast_1s1x64d"],
    "ResNet": [
        "ResNet18", "ResNet18_vd", "ResNet34", "ResNet34_vd", "ResNet50",
        "ResNet50_vc", "ResNet50_vd", "ResNet50_vd_v2", "ResNet101",
        "ResNet101_vd", "ResNet152", "ResNet152_vd", "ResNet200_vd",
        "ResNet34_vd_ssld", "ResNet50_vd_ssld", "ResNet50_vd_ssld_v2",
        "ResNet101_vd_ssld", "Fix_ResNet50_vd_ssld_v2", "ResNet50_ACNet_deploy"
    ],
    "ResNeXt": [
        "ResNeXt50_32x4d", "ResNeXt50_vd_32x4d", "ResNeXt50_64x4d",
        "ResNeXt50_vd_64x4d", "ResNeXt101_32x4d", "ResNeXt101_vd_32x4d",
        "ResNeXt101_32x8d_wsl", "ResNeXt101_32x16d_wsl",
        "ResNeXt101_32x32d_wsl", "ResNeXt101_32x48d_wsl",
        "Fix_ResNeXt101_32x48d_wsl", "ResNeXt101_64x4d", "ResNeXt101_vd_64x4d",
        "ResNeXt152_32x4d", "ResNeXt152_vd_32x4d", "ResNeXt152_64x4d",
        "ResNeXt152_vd_64x4d"
    ],
150 151
    "ReXNet":
    ["ReXNet_1_0", "ReXNet_1_3", "ReXNet_1_5", "ReXNet_2_0", "ReXNet_3_0"],
T
Tingquan Gao 已提交
152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172
    "SENet": [
        "SENet154_vd", "SE_HRNet_W64_C_ssld", "SE_ResNet18_vd",
        "SE_ResNet34_vd", "SE_ResNet50_vd", "SE_ResNeXt50_32x4d",
        "SE_ResNeXt50_vd_32x4d", "SE_ResNeXt101_32x4d"
    ],
    "ShuffleNetV2": [
        "ShuffleNetV2_swish", "ShuffleNetV2_x0_25", "ShuffleNetV2_x0_33",
        "ShuffleNetV2_x0_5", "ShuffleNetV2_x1_0", "ShuffleNetV2_x1_5",
        "ShuffleNetV2_x2_0"
    ],
    "SqueezeNet": ["SqueezeNet1_0", "SqueezeNet1_1"],
    "SwinTransformer": [
        "SwinTransformer_large_patch4_window7_224_22kto1k",
        "SwinTransformer_large_patch4_window12_384_22kto1k",
        "SwinTransformer_base_patch4_window7_224_22kto1k",
        "SwinTransformer_base_patch4_window12_384_22kto1k",
        "SwinTransformer_base_patch4_window12_384",
        "SwinTransformer_base_patch4_window7_224",
        "SwinTransformer_small_patch4_window7_224",
        "SwinTransformer_tiny_patch4_window7_224"
    ],
G
gaotingquan 已提交
173 174 175 176
    "Twins": [
        "pcpvt_small", "pcpvt_base", "pcpvt_large", "alt_gvt_small",
        "alt_gvt_base", "alt_gvt_large"
    ],
G
gaotingquan 已提交
177
    "TNT": ["TNT_small"],
T
Tingquan Gao 已提交
178 179 180 181 182 183 184 185 186 187
    "VGG": ["VGG11", "VGG13", "VGG16", "VGG19"],
    "VisionTransformer": [
        "ViT_base_patch16_224", "ViT_base_patch16_384", "ViT_base_patch32_384",
        "ViT_large_patch16_224", "ViT_large_patch16_384",
        "ViT_large_patch32_384", "ViT_small_patch16_224"
    ],
    "Xception": [
        "Xception41", "Xception41_deeplab", "Xception65", "Xception65_deeplab",
        "Xception71"
    ]
C
chenziheng 已提交
188 189
}

G
gaotingquan 已提交
190
PULC_MODEL_BASE_DOWNLOAD_URL = "https://paddleclas.bj.bcebos.com/models/PULC/inference/{}_infer.tar"
191
PULC_MODELS = [
G
gaotingquan 已提交
192 193
    "car_exists", "language_classification", "person_attribute",
    "person_exists", "safety_helmet", "text_image_orientation",
C
cuicheng01 已提交
194 195
    "textline_orientation", "traffic_sign", "vehicle_attribute",
    "table_attribute"
196 197
]

C
chenziheng 已提交
198

T
Tingquan Gao 已提交
199 200
class ImageTypeError(Exception):
    """ImageTypeError.
T
Tingquan Gao 已提交
201 202
    """

T
Tingquan Gao 已提交
203
    def __init__(self, message=""):
T
Tingquan Gao 已提交
204 205 206
        super().__init__(message)


T
Tingquan Gao 已提交
207 208 209 210 211 212 213 214
class InputModelError(Exception):
    """InputModelError.
    """

    def __init__(self, message=""):
        super().__init__(message)


215 216 217
def init_config(model_type, model_name, inference_model_dir, **kwargs):

    cfg_path = f"deploy/configs/PULC/{model_name}/inference_{model_name}.yaml" if model_type == "pulc" else "deploy/configs/inference_cls.yaml"
218
    __dir__ = os.path.dirname(__file__)
G
gaotingquan 已提交
219
    cfg_path = os.path.join(__dir__, cfg_path)
220 221 222 223 224 225
    cfg = config.get_config(cfg_path, show=False)

    cfg.Global.inference_model_dir = inference_model_dir

    if "batch_size" in kwargs and kwargs["batch_size"]:
        cfg.Global.batch_size = kwargs["batch_size"]
226

227 228
    if "use_gpu" in kwargs and kwargs["use_gpu"]:
        cfg.Global.use_gpu = kwargs["use_gpu"]
229 230 231 232
    if cfg.Global.use_gpu and not paddle.device.is_compiled_with_cuda():
        msg = "The current running environment does not support the use of GPU. CPU has been used instead."
        logger.warning(msg)
        cfg.Global.use_gpu = False
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

    if "infer_imgs" in kwargs and kwargs["infer_imgs"]:
        cfg.Global.infer_imgs = kwargs["infer_imgs"]
    if "enable_mkldnn" in kwargs and kwargs["enable_mkldnn"]:
        cfg.Global.enable_mkldnn = kwargs["enable_mkldnn"]
    if "cpu_num_threads" in kwargs and kwargs["cpu_num_threads"]:
        cfg.Global.cpu_num_threads = kwargs["cpu_num_threads"]
    if "use_fp16" in kwargs and kwargs["use_fp16"]:
        cfg.Global.use_fp16 = kwargs["use_fp16"]
    if "use_tensorrt" in kwargs and kwargs["use_tensorrt"]:
        cfg.Global.use_tensorrt = kwargs["use_tensorrt"]
    if "gpu_mem" in kwargs and kwargs["gpu_mem"]:
        cfg.Global.gpu_mem = kwargs["gpu_mem"]
    if "resize_short" in kwargs and kwargs["resize_short"]:
        cfg.PreProcess.transform_ops[0]["ResizeImage"][
            "resize_short"] = kwargs["resize_short"]
    if "crop_size" in kwargs and kwargs["crop_size"]:
        cfg.PreProcess.transform_ops[1]["CropImage"]["size"] = kwargs[
            "crop_size"]

    # TODO(gaotingquan): not robust
    if "thresh" in kwargs and kwargs[
            "thresh"] and "ThreshOutput" in cfg.PostProcess:
        cfg.PostProcess.ThreshOutput.thresh = kwargs["thresh"]
    if "Topk" in cfg.PostProcess:
        if "topk" in kwargs and kwargs["topk"]:
            cfg.PostProcess.Topk.topk = kwargs["topk"]
        if "class_id_map_file" in kwargs and kwargs["class_id_map_file"]:
            cfg.PostProcess.Topk.class_id_map_file = kwargs[
T
Tingquan Gao 已提交
262
                "class_id_map_file"]
263
        else:
G
gaotingquan 已提交
264
            class_id_map_file_path = os.path.relpath(
265
                cfg.PostProcess.Topk.class_id_map_file, "../")
G
gaotingquan 已提交
266 267
            cfg.PostProcess.Topk.class_id_map_file = os.path.join(
                __dir__, class_id_map_file_path)
268 269 270 271 272 273 274
    if "VehicleAttribute" in cfg.PostProcess:
        if "color_threshold" in kwargs and kwargs["color_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "color_threshold"]
        if "type_threshold" in kwargs and kwargs["type_threshold"]:
            cfg.PostProcess.VehicleAttribute.type_threshold = kwargs[
                "type_threshold"]
C
cuicheng01 已提交
275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293
    if "TableAttribute" in cfg.PostProcess:
        if "source_threshold" in kwargs and kwargs["source_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "source_threshold"]
        if "number_threshold" in kwargs and kwargs["number_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "number_threshold"]
        if "color_threshold" in kwargs and kwargs["color_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "color_threshold"]
        if "clarity_threshold" in kwargs and kwargs["clarity_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "clarity_threshold"]
        if "obstruction_threshold" in kwargs and kwargs["obstruction_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "obstruction_threshold"]
        if "angle_threshold" in kwargs and kwargs["angle_threshold"]:
            cfg.PostProcess.VehicleAttribute.color_threshold = kwargs[
                "angle_threshold"]
294 295
    if "save_dir" in kwargs and kwargs["save_dir"]:
        cfg.PostProcess.SavePreLabel.save_dir = kwargs["save_dir"]
T
Tingquan Gao 已提交
296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316

    return cfg


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

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--infer_imgs",
        type=str,
        required=True,
        help="The image(s) to be predicted.")
    parser.add_argument(
        "--model_name", type=str, help="The model name to be used.")
    parser.add_argument(
        "--inference_model_dir",
        type=str,
        help="The directory of model files. Valid when model_name not specifed."
    )
317 318 319 320 321
    parser.add_argument("--use_gpu", type=str2bool, help="Whether use GPU.")
    parser.add_argument(
        "--gpu_mem",
        type=int,
        help="The memory size of GPU allocated to predict.")
T
Tingquan Gao 已提交
322 323 324 325 326
    parser.add_argument(
        "--enable_mkldnn",
        type=str2bool,
        help="Whether use MKLDNN. Valid when use_gpu is False")
    parser.add_argument(
327 328 329 330 331 332
        "--cpu_num_threads",
        type=int,
        help="The threads number when predicting on CPU.")
    parser.add_argument(
        "--use_tensorrt",
        type=str2bool,
G
gaotingquan 已提交
333
        help="Whether use TensorRT to accelerate.")
334 335
    parser.add_argument(
        "--use_fp16", type=str2bool, help="Whether use FP16 to predict.")
336
    parser.add_argument("--batch_size", type=int, help="Batch size.")
T
Tingquan Gao 已提交
337 338 339
    parser.add_argument(
        "--topk",
        type=int,
340 341
        help="Return topk score(s) and corresponding results when Topk postprocess is used."
    )
T
Tingquan Gao 已提交
342 343 344 345
    parser.add_argument(
        "--class_id_map_file",
        type=str,
        help="The path of file that map class_id and label.")
346 347 348 349 350 351
    parser.add_argument(
        "--threshold",
        type=float,
        help="The threshold of ThreshOutput when postprocess is used.")
    parser.add_argument("--color_threshold", type=float, help="")
    parser.add_argument("--type_threshold", type=float, help="")
T
Tingquan Gao 已提交
352 353 354 355
    parser.add_argument(
        "--save_dir",
        type=str,
        help="The directory to save prediction results as pre-label.")
G
gaotingquan 已提交
356
    parser.add_argument(
357 358
        "--resize_short", type=int, help="Resize according to short size.")
    parser.add_argument("--crop_size", type=int, help="Centor crop size.")
T
Tingquan Gao 已提交
359 360 361

    args = parser.parse_args()
    return vars(args)
T
Tingquan Gao 已提交
362 363


T
Tingquan Gao 已提交
364
def print_info():
T
Tingquan Gao 已提交
365 366
    """Print list of supported models in formatted.
    """
367 368
    imn_table = PrettyTable(["IMN Model Series", "Model Name"])
    pulc_table = PrettyTable(["PULC Models"])
T
Tingquan Gao 已提交
369 370
    try:
        sz = os.get_terminal_size()
371 372 373
        total_width = sz.columns
        first_width = 30
        second_width = total_width - first_width if total_width > 50 else 10
T
Tingquan Gao 已提交
374
    except OSError:
375
        total_width = 100
376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396
        second_width = 100
    for series in IMN_MODEL_SERIES:
        names = textwrap.fill(
            "  ".join(IMN_MODEL_SERIES[series]), width=second_width)
        imn_table.add_row([series, names])

    table_width = len(str(imn_table).split("\n")[0])
    pulc_table.add_row([
        textwrap.fill(
            "  ".join(PULC_MODELS), width=total_width).center(table_width - 4)
    ])

    print("{}".format("-" * table_width))
    print("Models supported by PaddleClas".center(table_width))
    print(imn_table)
    print(pulc_table)
    print("Powered by PaddlePaddle!".rjust(table_width))
    print("{}".format("-" * table_width))


def get_imn_model_names():
T
Tingquan Gao 已提交
397 398
    """Get the model names list.
    """
T
Tingquan Gao 已提交
399
    model_names = []
400 401
    for series in IMN_MODEL_SERIES:
        model_names += (IMN_MODEL_SERIES[series])
T
Tingquan Gao 已提交
402 403 404
    return model_names


405
def similar_model_names(name="", names=[], thresh=0.1, topk=5):
T
Tingquan Gao 已提交
406
    """Find the most similar topk model names.
T
Tingquan Gao 已提交
407 408 409
    """
    scores = []
    for idx, n in enumerate(names):
T
Tingquan Gao 已提交
410
        if n.startswith("__"):
T
Tingquan Gao 已提交
411 412 413 414 415 416 417 418 419
            continue
        score = SequenceMatcher(None, n.lower(), name.lower()).quick_ratio()
        if score > thresh:
            scores.append((idx, score))
    scores.sort(key=lambda x: x[1], reverse=True)
    similar_names = [names[s[0]] for s in scores[:min(topk, len(scores))]]
    return similar_names


C
chenziheng 已提交
420
def download_with_progressbar(url, save_path):
T
Tingquan Gao 已提交
421 422 423 424
    """Download from url with progressbar.
    """
    if os.path.isfile(save_path):
        os.remove(save_path)
C
chenziheng 已提交
425
    response = requests.get(url, stream=True)
T
Tingquan Gao 已提交
426
    total_size_in_bytes = int(response.headers.get("content-length", 0))
C
chenziheng 已提交
427
    block_size = 1024  # 1 Kibibyte
T
Tingquan Gao 已提交
428 429
    progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
    with open(save_path, "wb") as file:
C
chenziheng 已提交
430 431 432 433
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)
    progress_bar.close()
T
Tingquan Gao 已提交
434 435
    if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes or not os.path.isfile(
            save_path):
T
Tingquan Gao 已提交
436
        raise Exception(
T
Tingquan Gao 已提交
437
            f"Something went wrong while downloading file from {url}")
C
chenziheng 已提交
438 439


440
def check_model_file(model_type, model_name):
441
    """Check the model files exist and download and untar when no exist.
T
Tingquan Gao 已提交
442
    """
443 444 445 446 447 448 449 450
    if model_type == "pulc":
        storage_directory = partial(os.path.join, BASE_INFERENCE_MODEL_DIR,
                                    "PULC", model_name)
        url = PULC_MODEL_BASE_DOWNLOAD_URL.format(model_name)
    else:
        storage_directory = partial(os.path.join, BASE_INFERENCE_MODEL_DIR,
                                    "IMN", model_name)
        url = IMN_MODEL_BASE_DOWNLOAD_URL.format(model_name)
T
Tingquan Gao 已提交
451

C
chenziheng 已提交
452
    tar_file_name_list = [
T
Tingquan Gao 已提交
453
        "inference.pdiparams", "inference.pdiparams.info", "inference.pdmodel"
C
chenziheng 已提交
454
    ]
T
Tingquan Gao 已提交
455 456 457 458 459
    model_file_path = storage_directory("inference.pdmodel")
    params_file_path = storage_directory("inference.pdiparams")
    if not os.path.exists(model_file_path) or not os.path.exists(
            params_file_path):
        tmp_path = storage_directory(url.split("/")[-1])
G
gaotingquan 已提交
460
        logger.info(f"download {url} to {tmp_path}")
T
Tingquan Gao 已提交
461
        os.makedirs(storage_directory(), exist_ok=True)
C
chenziheng 已提交
462
        download_with_progressbar(url, tmp_path)
T
Tingquan Gao 已提交
463
        with tarfile.open(tmp_path, "r") as tarObj:
C
chenziheng 已提交
464 465 466 467 468 469 470 471
            for member in tarObj.getmembers():
                filename = None
                for tar_file_name in tar_file_name_list:
                    if tar_file_name in member.name:
                        filename = tar_file_name
                if filename is None:
                    continue
                file = tarObj.extractfile(member)
T
Tingquan Gao 已提交
472
                with open(storage_directory(filename), "wb") as f:
C
chenziheng 已提交
473 474
                    f.write(file.read())
        os.remove(tmp_path)
T
Tingquan Gao 已提交
475 476 477 478 479
    if not os.path.exists(model_file_path) or not os.path.exists(
            params_file_path):
        raise Exception(
            f"Something went wrong while praparing the model[{model_name}] files!"
        )
C
chenziheng 已提交
480

T
Tingquan Gao 已提交
481
    return storage_directory()
C
chenziheng 已提交
482

T
Tingquan Gao 已提交
483

C
chenziheng 已提交
484
class PaddleClas(object):
T
Tingquan Gao 已提交
485 486 487 488 489 490
    """PaddleClas.
    """

    def __init__(self,
                 model_name: str=None,
                 inference_model_dir: str=None,
T
Tingquan Gao 已提交
491
                 **kwargs):
T
Tingquan Gao 已提交
492
        """Init PaddleClas with config.
T
Tingquan Gao 已提交
493

T
Tingquan Gao 已提交
494
        Args:
495 496 497 498 499
            model_name (str, optional): The model name supported by PaddleClas. If specified, override config. Defaults to None.
            inference_model_dir (str, optional): The directory that contained model file and params file to be used. If specified, override config. Defaults to None.
            use_gpu (bool, optional): Whether use GPU. If specified, override config. Defaults to True.
            batch_size (int, optional): The batch size to pridict. If specified, override config. Defaults to 1.
            topk (int, optional): Return the top k prediction results with the highest score. Defaults to 5.
T
Tingquan Gao 已提交
500 501
        """
        super().__init__()
502

503 504 505 506 507
        self.model_type, inference_model_dir = self._check_input_model(
            model_name, inference_model_dir)
        self._config = init_config(self.model_type, model_name,
                                   inference_model_dir, **kwargs)

T
Tingquan Gao 已提交
508 509 510 511
        self.cls_predictor = ClsPredictor(self._config)

    def get_config(self):
        """Get the config.
C
chenziheng 已提交
512
        """
T
Tingquan Gao 已提交
513 514
        return self._config

515
    def _check_input_model(self, model_name, inference_model_dir):
T
Tingquan Gao 已提交
516 517
        """Check input model name or model files.
        """
518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535
        all_imn_model_names = get_imn_model_names()
        all_pulc_model_names = PULC_MODELS

        if model_name:
            if model_name in all_imn_model_names:
                inference_model_dir = check_model_file("imn", model_name)
                return "imn", inference_model_dir
            elif model_name in all_pulc_model_names:
                inference_model_dir = check_model_file("pulc", model_name)
                return "pulc", inference_model_dir
            else:
                similar_imn_names = similar_model_names(model_name,
                                                        all_imn_model_names)
                similar_pulc_names = similar_model_names(model_name,
                                                         all_pulc_model_names)
                similar_names_str = ", ".join(similar_imn_names +
                                              similar_pulc_names)
                err = f"{model_name} is not provided by PaddleClas. \nMaybe you want the : [{similar_names_str}]. \nIf you want to use your own model, please specify inference_model_dir!"
T
Tingquan Gao 已提交
536
                raise InputModelError(err)
537
        elif inference_model_dir:
T
Tingquan Gao 已提交
538 539 540 541 542 543 544 545
            model_file_path = os.path.join(inference_model_dir,
                                           "inference.pdmodel")
            params_file_path = os.path.join(inference_model_dir,
                                            "inference.pdiparams")
            if not os.path.isfile(model_file_path) or not os.path.isfile(
                    params_file_path):
                err = f"There is no model file or params file in this directory: {inference_model_dir}"
                raise InputModelError(err)
546
            return "custom", inference_model_dir
T
Tingquan Gao 已提交
547
        else:
T
Tingquan Gao 已提交
548
            err = f"Please specify the model name supported by PaddleClas or directory contained model files(inference.pdmodel, inference.pdiparams)."
T
Tingquan Gao 已提交
549
            raise InputModelError(err)
550
        return None
T
Tingquan Gao 已提交
551

552 553
    def predict(self, input_data: Union[str, np.array],
                print_pred: bool=False) -> Generator[list, None, None]:
T
Tingquan Gao 已提交
554 555
        """Predict input_data.

C
chenziheng 已提交
556
        Args:
G
gaotingquan 已提交
557
            input_data (Union[str, np.array]):
558 559
                When the type is str, it is the path of image, or the directory containing images, or the URL of image from Internet.
                When the type is np.array, it is the image data whose channel order is RGB.
G
gaotingquan 已提交
560
            print_pred (bool, optional): Whether print the prediction result. Defaults to False.
T
Tingquan Gao 已提交
561 562 563 564 565

        Raises:
            ImageTypeError: Illegal input_data.

        Yields:
G
gaotingquan 已提交
566 567 568
            Generator[list, None, None]:
                The prediction result(s) of input_data by batch_size. For every one image,
                prediction result(s) is zipped as a dict, that includs topk "class_ids", "scores" and "label_names".
G
gaotingquan 已提交
569
                The format of batch prediction result(s) is as follow: [{"class_ids": [...], "scores": [...], "label_names": [...]}, ...]
C
chenziheng 已提交
570
        """
571

T
Tingquan Gao 已提交
572
        if isinstance(input_data, np.ndarray):
G
gaotingquan 已提交
573
            yield self.cls_predictor.predict(input_data)
T
Tingquan Gao 已提交
574
        elif isinstance(input_data, str):
T
Tingquan Gao 已提交
575
            if input_data.startswith("http") or input_data.startswith("https"):
T
Tingquan Gao 已提交
576 577 578 579 580
                image_storage_dir = partial(os.path.join, BASE_IMAGES_DIR)
                if not os.path.exists(image_storage_dir()):
                    os.makedirs(image_storage_dir())
                image_save_path = image_storage_dir("tmp.jpg")
                download_with_progressbar(input_data, image_save_path)
G
gaotingquan 已提交
581
                logger.info(
T
Tingquan Gao 已提交
582 583
                    f"Image to be predicted from Internet: {input_data}, has been saved to: {image_save_path}"
                )
584
                input_data = image_save_path
T
Tingquan Gao 已提交
585 586 587 588 589
            image_list = get_image_list(input_data)

            batch_size = self._config.Global.get("batch_size", 1)

            img_list = []
T
Tingquan Gao 已提交
590 591
            img_path_list = []
            cnt = 0
592
            for idx_img, img_path in enumerate(image_list):
T
Tingquan Gao 已提交
593 594
                img = cv2.imread(img_path)
                if img is None:
G
gaotingquan 已提交
595
                    logger.warning(
T
Tingquan Gao 已提交
596 597
                        f"Image file failed to read and has been skipped. The path: {img_path}"
                    )
T
Tingquan Gao 已提交
598
                    continue
599
                img = img[:, :, ::-1]
T
Tingquan Gao 已提交
600 601 602 603
                img_list.append(img)
                img_path_list.append(img_path)
                cnt += 1

604
                if cnt % batch_size == 0 or (idx_img + 1) == len(image_list):
G
gaotingquan 已提交
605 606
                    preds = self.cls_predictor.predict(img_list)

607 608 609 610
                    if preds:
                        for idx_pred, pred in enumerate(preds):
                            pred["filename"] = img_path_list[idx_pred]
                            if print_pred:
G
gaotingquan 已提交
611
                                logger.info(", ".join(
612
                                    [f"{k}: {pred[k]}" for k in pred]))
T
Tingquan Gao 已提交
613

T
Tingquan Gao 已提交
614
                    img_list = []
T
Tingquan Gao 已提交
615
                    img_path_list = []
T
Tingquan Gao 已提交
616
                    yield preds
C
chenziheng 已提交
617
        else:
T
Tingquan Gao 已提交
618 619 620
            err = "Please input legal image! The type of image supported by PaddleClas are: NumPy.ndarray and string of local path or Ineternet URL"
            raise ImageTypeError(err)
        return
C
chenziheng 已提交
621 622


T
Tingquan Gao 已提交
623
# for CLI
C
chenziheng 已提交
624
def main():
T
Tingquan Gao 已提交
625 626
    """Function API used for commad line.
    """
627
    print_info()
T
Tingquan Gao 已提交
628
    cfg = args_cfg()
T
Tingquan Gao 已提交
629 630 631 632
    clas_engine = PaddleClas(**cfg)
    res = clas_engine.predict(cfg["infer_imgs"], print_pred=True)
    for _ in res:
        pass
G
gaotingquan 已提交
633
    logger.info("Predict complete!")
T
Tingquan Gao 已提交
634
    return
C
chenziheng 已提交
635 636


T
Tingquan Gao 已提交
637
if __name__ == "__main__":
C
chenziheng 已提交
638
    main()