paddleclas.py 23.9 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 16 17
#
# 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
import sys
__dir__ = os.path.dirname(__file__)
T
Tingquan Gao 已提交
18 19 20
sys.path.append(os.path.join(__dir__, ""))
sys.path.append(os.path.join(__dir__, "deploy"))

21
from typing import Union, Generator
T
Tingquan Gao 已提交
22 23
import argparse
import shutil
T
Tingquan Gao 已提交
24
import textwrap
T
Tingquan Gao 已提交
25 26 27
import tarfile
import requests
from functools import partial
T
Tingquan Gao 已提交
28
from difflib import SequenceMatcher
C
chenziheng 已提交
29 30 31 32

import cv2
import numpy as np
from tqdm import tqdm
T
Tingquan Gao 已提交
33
from prettytable import PrettyTable
34
import paddle
T
Tingquan Gao 已提交
35 36 37 38 39

from deploy.python.predict_cls import ClsPredictor
from deploy.utils.get_image_list import get_image_list
from deploy.utils import config

40
import ppcls.arch.backbone as backbone
G
gaotingquan 已提交
41
from ppcls.utils import logger
T
Tingquan Gao 已提交
42

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

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

C
chenziheng 已提交
48
BASE_DIR = os.path.expanduser("~/.paddleclas/")
T
Tingquan Gao 已提交
49 50
BASE_INFERENCE_MODEL_DIR = os.path.join(BASE_DIR, "inference_model")
BASE_IMAGES_DIR = os.path.join(BASE_DIR, "images")
51 52
IMN_MODEL_BASE_DOWNLOAD_URL = "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/inference/{}_infer.tar"
IMN_MODEL_SERIES = {
T
Tingquan Gao 已提交
53 54 55
    "AlexNet": ["AlexNet"],
    "DarkNet": ["DarkNet53"],
    "DeiT": [
T
Tingquan Gao 已提交
56 57 58 59
        "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 已提交
60 61 62 63 64
    ],
    "DenseNet": [
        "DenseNet121", "DenseNet161", "DenseNet169", "DenseNet201",
        "DenseNet264"
    ],
65 66 67 68
    "DLA": [
        "DLA46_c", "DLA60x_c", "DLA34", "DLA60", "DLA60x", "DLA102", "DLA102x",
        "DLA102x2", "DLA169"
    ],
T
Tingquan Gao 已提交
69 70 71 72 73 74
    "DPN": ["DPN68", "DPN92", "DPN98", "DPN107", "DPN131"],
    "EfficientNet": [
        "EfficientNetB0", "EfficientNetB0_small", "EfficientNetB1",
        "EfficientNetB2", "EfficientNetB3", "EfficientNetB4", "EfficientNetB5",
        "EfficientNetB6", "EfficientNetB7"
    ],
G
gaotingquan 已提交
75
    "ESNet": ["ESNet_x0_25", "ESNet_x0_5", "ESNet_x0_75", "ESNet_x1_0"],
T
Tingquan Gao 已提交
76 77
    "GhostNet":
    ["GhostNet_x0_5", "GhostNet_x1_0", "GhostNet_x1_3", "GhostNet_x1_3_ssld"],
78
    "HarDNet": ["HarDNet39_ds", "HarDNet68_ds", "HarDNet68", "HarDNet85"],
T
Tingquan Gao 已提交
79 80 81 82 83 84
    "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 已提交
85
    "MixNet": ["MixNet_S", "MixNet_M", "MixNet_L"],
T
Tingquan Gao 已提交
86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102
    "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 已提交
103 104 105 106 107 108
    "PPHGNet": [
        "PPHGNet_tiny",
        "PPHGNet_small",
        "PPHGNet_tiny_ssld",
        "PPHGNet_small_ssld",
    ],
G
gaotingquan 已提交
109 110 111 112
    "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 已提交
113
    "PPLCNetV2": ["PPLCNetV2_base"],
114
    "RedNet": ["RedNet26", "RedNet38", "RedNet50", "RedNet101", "RedNet152"],
T
Tingquan Gao 已提交
115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138
    "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"
    ],
139 140
    "ReXNet":
    ["ReXNet_1_0", "ReXNet_1_3", "ReXNet_1_5", "ReXNet_2_0", "ReXNet_3_0"],
T
Tingquan Gao 已提交
141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161
    "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 已提交
162 163 164 165
    "Twins": [
        "pcpvt_small", "pcpvt_base", "pcpvt_large", "alt_gvt_small",
        "alt_gvt_base", "alt_gvt_large"
    ],
T
Tingquan Gao 已提交
166 167 168 169 170 171 172 173 174 175
    "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 已提交
176 177
}

G
gaotingquan 已提交
178
PULC_MODEL_BASE_DOWNLOAD_URL = "https://paddleclas.bj.bcebos.com/models/PULC/inference/{}_infer.tar"
179
PULC_MODELS = [
G
gaotingquan 已提交
180 181 182
    "car_exists", "language_classification", "person_attribute",
    "person_exists", "safety_helmet", "text_image_orientation",
    "textline_orientation", "traffic_sign", "vehicle_attribute"
183 184
]

C
chenziheng 已提交
185

T
Tingquan Gao 已提交
186 187
class ImageTypeError(Exception):
    """ImageTypeError.
T
Tingquan Gao 已提交
188 189
    """

T
Tingquan Gao 已提交
190
    def __init__(self, message=""):
T
Tingquan Gao 已提交
191 192 193
        super().__init__(message)


T
Tingquan Gao 已提交
194 195 196 197 198 199 200 201
class InputModelError(Exception):
    """InputModelError.
    """

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


202 203 204
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"
G
gaotingquan 已提交
205
    cfg_path = os.path.join(__dir__, cfg_path)
206 207 208 209 210 211
    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"]
212

213 214
    if "use_gpu" in kwargs and kwargs["use_gpu"]:
        cfg.Global.use_gpu = kwargs["use_gpu"]
215 216 217 218
    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
219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247

    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 已提交
248
                "class_id_map_file"]
249
        else:
G
gaotingquan 已提交
250
            class_id_map_file_path = os.path.relpath(
251
                cfg.PostProcess.Topk.class_id_map_file, "../")
G
gaotingquan 已提交
252 253
            cfg.PostProcess.Topk.class_id_map_file = os.path.join(
                __dir__, class_id_map_file_path)
254 255 256 257 258 259 260 261 262 263
    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"]

    if "save_dir" in kwargs and kwargs["save_dir"]:
        cfg.PostProcess.SavePreLabel.save_dir = kwargs["save_dir"]
T
Tingquan Gao 已提交
264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284

    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."
    )
285 286 287 288 289
    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 已提交
290 291 292 293 294
    parser.add_argument(
        "--enable_mkldnn",
        type=str2bool,
        help="Whether use MKLDNN. Valid when use_gpu is False")
    parser.add_argument(
295 296 297 298 299 300
        "--cpu_num_threads",
        type=int,
        help="The threads number when predicting on CPU.")
    parser.add_argument(
        "--use_tensorrt",
        type=str2bool,
G
gaotingquan 已提交
301
        help="Whether use TensorRT to accelerate.")
302 303
    parser.add_argument(
        "--use_fp16", type=str2bool, help="Whether use FP16 to predict.")
304
    parser.add_argument("--batch_size", type=int, help="Batch size.")
T
Tingquan Gao 已提交
305 306 307
    parser.add_argument(
        "--topk",
        type=int,
308 309
        help="Return topk score(s) and corresponding results when Topk postprocess is used."
    )
T
Tingquan Gao 已提交
310 311 312 313
    parser.add_argument(
        "--class_id_map_file",
        type=str,
        help="The path of file that map class_id and label.")
314 315 316 317 318 319
    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 已提交
320 321 322 323
    parser.add_argument(
        "--save_dir",
        type=str,
        help="The directory to save prediction results as pre-label.")
G
gaotingquan 已提交
324
    parser.add_argument(
325 326
        "--resize_short", type=int, help="Resize according to short size.")
    parser.add_argument("--crop_size", type=int, help="Centor crop size.")
T
Tingquan Gao 已提交
327 328 329

    args = parser.parse_args()
    return vars(args)
T
Tingquan Gao 已提交
330 331


T
Tingquan Gao 已提交
332
def print_info():
T
Tingquan Gao 已提交
333 334
    """Print list of supported models in formatted.
    """
335 336
    imn_table = PrettyTable(["IMN Model Series", "Model Name"])
    pulc_table = PrettyTable(["PULC Models"])
T
Tingquan Gao 已提交
337 338
    try:
        sz = os.get_terminal_size()
339 340 341
        total_width = sz.columns
        first_width = 30
        second_width = total_width - first_width if total_width > 50 else 10
T
Tingquan Gao 已提交
342
    except OSError:
343
        total_width = 100
344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364
        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 已提交
365 366
    """Get the model names list.
    """
T
Tingquan Gao 已提交
367
    model_names = []
368 369
    for series in IMN_MODEL_SERIES:
        model_names += (IMN_MODEL_SERIES[series])
T
Tingquan Gao 已提交
370 371 372
    return model_names


373
def similar_model_names(name="", names=[], thresh=0.1, topk=5):
T
Tingquan Gao 已提交
374
    """Find the most similar topk model names.
T
Tingquan Gao 已提交
375 376 377
    """
    scores = []
    for idx, n in enumerate(names):
T
Tingquan Gao 已提交
378
        if n.startswith("__"):
T
Tingquan Gao 已提交
379 380 381 382 383 384 385 386 387
            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 已提交
388
def download_with_progressbar(url, save_path):
T
Tingquan Gao 已提交
389 390 391 392
    """Download from url with progressbar.
    """
    if os.path.isfile(save_path):
        os.remove(save_path)
C
chenziheng 已提交
393
    response = requests.get(url, stream=True)
T
Tingquan Gao 已提交
394
    total_size_in_bytes = int(response.headers.get("content-length", 0))
C
chenziheng 已提交
395
    block_size = 1024  # 1 Kibibyte
T
Tingquan Gao 已提交
396 397
    progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
    with open(save_path, "wb") as file:
C
chenziheng 已提交
398 399 400 401
        for data in response.iter_content(block_size):
            progress_bar.update(len(data))
            file.write(data)
    progress_bar.close()
T
Tingquan Gao 已提交
402 403
    if total_size_in_bytes == 0 or progress_bar.n != total_size_in_bytes or not os.path.isfile(
            save_path):
T
Tingquan Gao 已提交
404
        raise Exception(
T
Tingquan Gao 已提交
405
            f"Something went wrong while downloading file from {url}")
C
chenziheng 已提交
406 407


408
def check_model_file(model_type, model_name):
409
    """Check the model files exist and download and untar when no exist.
T
Tingquan Gao 已提交
410
    """
411 412 413 414 415 416 417 418
    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 已提交
419

C
chenziheng 已提交
420
    tar_file_name_list = [
T
Tingquan Gao 已提交
421
        "inference.pdiparams", "inference.pdiparams.info", "inference.pdmodel"
C
chenziheng 已提交
422
    ]
T
Tingquan Gao 已提交
423 424 425 426 427
    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 已提交
428
        logger.info(f"download {url} to {tmp_path}")
T
Tingquan Gao 已提交
429
        os.makedirs(storage_directory(), exist_ok=True)
C
chenziheng 已提交
430
        download_with_progressbar(url, tmp_path)
T
Tingquan Gao 已提交
431
        with tarfile.open(tmp_path, "r") as tarObj:
C
chenziheng 已提交
432 433 434 435 436 437 438 439
            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 已提交
440
                with open(storage_directory(filename), "wb") as f:
C
chenziheng 已提交
441 442
                    f.write(file.read())
        os.remove(tmp_path)
T
Tingquan Gao 已提交
443 444 445 446 447
    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 已提交
448

T
Tingquan Gao 已提交
449
    return storage_directory()
C
chenziheng 已提交
450

T
Tingquan Gao 已提交
451

C
chenziheng 已提交
452
class PaddleClas(object):
T
Tingquan Gao 已提交
453 454 455
    """PaddleClas.
    """

456 457 458
    if not os.environ.get('ppcls', False):
        os.environ.setdefault('ppcls', 'True')
        print_info()
C
chenziheng 已提交
459

T
Tingquan Gao 已提交
460 461 462
    def __init__(self,
                 model_name: str=None,
                 inference_model_dir: str=None,
T
Tingquan Gao 已提交
463
                 **kwargs):
T
Tingquan Gao 已提交
464
        """Init PaddleClas with config.
T
Tingquan Gao 已提交
465

T
Tingquan Gao 已提交
466
        Args:
467 468 469 470 471
            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 已提交
472 473
        """
        super().__init__()
474 475 476 477 478
        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 已提交
479 480 481 482
        self.cls_predictor = ClsPredictor(self._config)

    def get_config(self):
        """Get the config.
C
chenziheng 已提交
483
        """
T
Tingquan Gao 已提交
484 485
        return self._config

486
    def _check_input_model(self, model_name, inference_model_dir):
T
Tingquan Gao 已提交
487 488
        """Check input model name or model files.
        """
489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506
        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 已提交
507
                raise InputModelError(err)
508
        elif inference_model_dir:
T
Tingquan Gao 已提交
509 510 511 512 513 514 515 516
            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)
517
            return "custom", inference_model_dir
T
Tingquan Gao 已提交
518
        else:
T
Tingquan Gao 已提交
519
            err = f"Please specify the model name supported by PaddleClas or directory contained model files(inference.pdmodel, inference.pdiparams)."
T
Tingquan Gao 已提交
520
            raise InputModelError(err)
521
        return None
T
Tingquan Gao 已提交
522

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

C
chenziheng 已提交
527
        Args:
G
gaotingquan 已提交
528
            input_data (Union[str, np.array]):
529 530
                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 已提交
531
            print_pred (bool, optional): Whether print the prediction result. Defaults to False.
T
Tingquan Gao 已提交
532 533 534 535 536

        Raises:
            ImageTypeError: Illegal input_data.

        Yields:
G
gaotingquan 已提交
537 538 539
            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 已提交
540
                The format of batch prediction result(s) is as follow: [{"class_ids": [...], "scores": [...], "label_names": [...]}, ...]
C
chenziheng 已提交
541
        """
542

T
Tingquan Gao 已提交
543
        if isinstance(input_data, np.ndarray):
G
gaotingquan 已提交
544
            yield self.cls_predictor.predict(input_data)
T
Tingquan Gao 已提交
545
        elif isinstance(input_data, str):
T
Tingquan Gao 已提交
546
            if input_data.startswith("http") or input_data.startswith("https"):
T
Tingquan Gao 已提交
547 548 549 550 551
                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 已提交
552
                logger.info(
T
Tingquan Gao 已提交
553 554
                    f"Image to be predicted from Internet: {input_data}, has been saved to: {image_save_path}"
                )
555
                input_data = image_save_path
T
Tingquan Gao 已提交
556 557 558 559 560
            image_list = get_image_list(input_data)

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

            img_list = []
T
Tingquan Gao 已提交
561 562
            img_path_list = []
            cnt = 0
563
            for idx_img, img_path in enumerate(image_list):
T
Tingquan Gao 已提交
564 565
                img = cv2.imread(img_path)
                if img is None:
G
gaotingquan 已提交
566
                    logger.warning(
T
Tingquan Gao 已提交
567 568
                        f"Image file failed to read and has been skipped. The path: {img_path}"
                    )
T
Tingquan Gao 已提交
569
                    continue
570
                img = img[:, :, ::-1]
T
Tingquan Gao 已提交
571 572 573 574
                img_list.append(img)
                img_path_list.append(img_path)
                cnt += 1

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

578 579 580 581
                    if preds:
                        for idx_pred, pred in enumerate(preds):
                            pred["filename"] = img_path_list[idx_pred]
                            if print_pred:
G
gaotingquan 已提交
582
                                logger.info(", ".join(
583
                                    [f"{k}: {pred[k]}" for k in pred]))
T
Tingquan Gao 已提交
584

T
Tingquan Gao 已提交
585
                    img_list = []
T
Tingquan Gao 已提交
586
                    img_path_list = []
T
Tingquan Gao 已提交
587
                    yield preds
C
chenziheng 已提交
588
        else:
T
Tingquan Gao 已提交
589 590 591
            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 已提交
592 593


T
Tingquan Gao 已提交
594
# for CLI
C
chenziheng 已提交
595
def main():
T
Tingquan Gao 已提交
596 597 598
    """Function API used for commad line.
    """
    cfg = args_cfg()
T
Tingquan Gao 已提交
599 600 601 602
    clas_engine = PaddleClas(**cfg)
    res = clas_engine.predict(cfg["infer_imgs"], print_pred=True)
    for _ in res:
        pass
G
gaotingquan 已提交
603
    logger.info("Predict complete!")
T
Tingquan Gao 已提交
604
    return
C
chenziheng 已提交
605 606


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