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
sys.path.append(os.path.join(__dir__, ""))

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

import cv2
import numpy as np
from tqdm import tqdm
T
Tingquan Gao 已提交
32
from prettytable import PrettyTable
33
import paddle
T
Tingquan Gao 已提交
34

35 36 37
import ppcls.arch.backbone as backbone
from ppcls.utils import logger

T
Tingquan Gao 已提交
38 39 40 41
from deploy.python.predict_cls import ClsPredictor
from deploy.utils.get_image_list import get_image_list
from deploy.utils import config

R
root 已提交
42
# for the PaddleClas Project to import
43 44
import deploy
import ppcls
T
Tingquan Gao 已提交
45

46
# for building model with loading pretrained weights from backbone
G
gaotingquan 已提交
47
logger.init_logger()
48

T
Tingquan Gao 已提交
49
__all__ = ["PaddleClas"]
T
Tingquan Gao 已提交
50

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

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

C
chenziheng 已提交
188

T
Tingquan Gao 已提交
189 190
class ImageTypeError(Exception):
    """ImageTypeError.
T
Tingquan Gao 已提交
191 192
    """

T
Tingquan Gao 已提交
193
    def __init__(self, message=""):
T
Tingquan Gao 已提交
194 195 196
        super().__init__(message)


T
Tingquan Gao 已提交
197 198 199 200 201 202 203 204
class InputModelError(Exception):
    """InputModelError.
    """

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


205 206 207
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 已提交
208
    cfg_path = os.path.join(__dir__, cfg_path)
209 210 211 212 213 214
    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"]
215

216 217
    if "use_gpu" in kwargs and kwargs["use_gpu"]:
        cfg.Global.use_gpu = kwargs["use_gpu"]
218 219 220 221
    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
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 248 249 250

    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 已提交
251
                "class_id_map_file"]
252
        else:
G
gaotingquan 已提交
253
            class_id_map_file_path = os.path.relpath(
254
                cfg.PostProcess.Topk.class_id_map_file, "../")
G
gaotingquan 已提交
255 256
            cfg.PostProcess.Topk.class_id_map_file = os.path.join(
                __dir__, class_id_map_file_path)
257 258 259 260 261 262 263 264 265 266
    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 已提交
267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287

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

    args = parser.parse_args()
    return vars(args)
T
Tingquan Gao 已提交
333 334


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


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


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

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

T
Tingquan Gao 已提交
452
    return storage_directory()
C
chenziheng 已提交
453

T
Tingquan Gao 已提交
454

C
chenziheng 已提交
455
class PaddleClas(object):
T
Tingquan Gao 已提交
456 457 458
    """PaddleClas.
    """

459 460 461
    if not os.environ.get('ppcls', False):
        os.environ.setdefault('ppcls', 'True')
        print_info()
C
chenziheng 已提交
462

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

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

    def get_config(self):
        """Get the config.
C
chenziheng 已提交
486
        """
T
Tingquan Gao 已提交
487 488
        return self._config

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

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

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

        Raises:
            ImageTypeError: Illegal input_data.

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

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

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

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

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

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

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


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


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