#! /usr/bin/env python
# -*- coding: utf-8 -*-
'''
Created on 2022/8/19 11:11:45

@author: LiuKuan
@copyright: Apache License, Version 2.0
'''
# ===============================================================================
'''
'''
# ===============================================================================

classify_list = [
    "alexnet",
    "googlenet",
    "inception_v3",
    "vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn",
    "convnext_base", "convnext_large", "convnext_small", "convnext_tiny",
    "densenet121", "densenet161","densenet169", "densenet201",
    "efficientnet_b0", "efficientnet_b1", "efficientnet_b2", "efficientnet_b3", "efficientnet_b4", "efficientnet_b5", "efficientnet_b6","efficientnet_b7",
    "efficientnet_v2_l", "efficientnet_v2_m", "efficientnet_v2_s",
    "mnasnet0_5", "mnasnet0_75", "mnasnet1_0", "mnasnet1_3",
    "mobilenet_v2", "mobilenet_v3_large", "mobilenet_v3_small",
    "regnet_x_16gf", "regnet_x_1_6gf", "regnet_x_32gf", "regnet_x_3_2gf", "regnet_x_400mf", "regnet_x_800mf", "regnet_x_8gf",
    "regnet_y_128gf", "regnet_y_16gf", "regnet_y_1_6gf", "regnet_y_32gf", "regnet_y_3_2gf", "regnet_y_400mf", "regnet_y_800mf", "regnet_y_8gf",
    "resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "resnext101_32x8d", "resnext101_64x4d", "resnext50_32x4d", "wide_resnet101_2", "wide_resnet50_2",
    "shufflenet_v2_x0_5", "shufflenet_v2_x1_0", "shufflenet_v2_x1_5", "shufflenet_v2_x2_0",
    "squeezenet1_0", "squeezenet1_1",
    "swin_b", "swin_s", "swin_t",
    "vit_b_16", "vit_b_32", "vit_h_14", "vit_l_16", "vit_l_32",
]

def get_classify(name, **kws):
    name = name.lower()
    if name == "alexnet":
        from torchvision.models import alexnet as Alg
    elif name == "convnext_base":
        from torchvision.models import convnext_base as Alg
    elif name == "convnext_large":
        from torchvision.models import convnext_large as Alg
    elif name == "convnext_small":
        from torchvision.models import convnext_small as Alg
    elif name == "convnext_tiny":
        from torchvision.models import convnext_tiny as Alg
    elif name == "densenet121":
        from torchvision.models import densenet121 as Alg
    elif name == "densenet161":
        from torchvision.models import densenet161 as Alg
    elif name == "densenet169":
        from torchvision.models import densenet169 as Alg
    elif name == "densenet201":
        from torchvision.models import densenet201 as Alg
    elif name == "efficientnet_b0":
        from torchvision.models import efficientnet_b0 as Alg
    elif name == "efficientnet_b1":
        from torchvision.models import efficientnet_b1 as Alg
    elif name == "efficientnet_b2":
        from torchvision.models import efficientnet_b2 as Alg
    elif name == "efficientnet_b3":
        from torchvision.models import efficientnet_b3 as Alg
    elif name == "efficientnet_b4":
        from torchvision.models import efficientnet_b4 as Alg
    elif name == "efficientnet_b5":
        from torchvision.models import efficientnet_b5 as Alg
    elif name == "efficientnet_b6":
        from torchvision.models import efficientnet_b6 as Alg
    elif name == "efficientnet_b7":
        from torchvision.models import efficientnet_b7 as Alg
    elif name == "efficientnet_v2_l":
        from torchvision.models import efficientnet_v2_l as Alg
    elif name == "efficientnet_v2_m":
        from torchvision.models import efficientnet_v2_m as Alg
    elif name == "efficientnet_v2_s":
        from torchvision.models import efficientnet_v2_s as Alg
    elif name == "googlenet":
        from torchvision.models import googlenet as Alg
    elif name == "inception_v3":
        from torchvision.models import inception_v3 as Alg
    elif name == "mnasnet0_5":
        from torchvision.models import mnasnet0_5 as Alg
    elif name == "mnasnet0_75":
        from torchvision.models import mnasnet0_75 as Alg
    elif name == "mnasnet1_0":
        from torchvision.models import mnasnet1_0 as Alg
    elif name == "mnasnet1_3":
        from torchvision.models import mnasnet1_3 as Alg
    elif name == "mobilenet_v2":
        from torchvision.models import mobilenet_v2 as Alg
    elif name == "mobilenet_v3_large":
        from torchvision.models import mobilenet_v3_large as Alg
    elif name == "mobilenet_v3_small":
        from torchvision.models import mobilenet_v3_small as Alg
    elif name == "regnet_x_16gf":
        from torchvision.models import regnet_x_16gf as Alg
    elif name == "regnet_x_1_6gf":
        from torchvision.models import regnet_x_1_6gf as Alg
    elif name == "regnet_x_32gf":
        from torchvision.models import regnet_x_32gf as Alg
    elif name == "regnet_x_3_2gf":
        from torchvision.models import regnet_x_3_2gf as Alg
    elif name == "regnet_x_400mf":
        from torchvision.models import regnet_x_400mf as Alg
    elif name == "regnet_x_800mf":
        from torchvision.models import regnet_x_800mf as Alg
    elif name == "regnet_x_8gf":
        from torchvision.models import regnet_x_8gf as Alg
    elif name == "regnet_y_128gf":
        from torchvision.models import regnet_y_128gf as Alg
    elif name == "regnet_y_16gf":
        from torchvision.models import regnet_y_16gf as Alg
    elif name == "regnet_y_1_6gf":
        from torchvision.models import regnet_y_1_6gf as Alg
    elif name == "regnet_y_32gf":
        from torchvision.models import regnet_y_32gf as Alg
    elif name == "regnet_y_3_2gf":
        from torchvision.models import regnet_y_3_2gf as Alg
    elif name == "regnet_y_400mf":
        from torchvision.models import regnet_y_400mf as Alg
    elif name == "regnet_y_800mf":
        from torchvision.models import regnet_y_800mf as Alg
    elif name == "regnet_y_8gf":
        from torchvision.models import regnet_y_8gf as Alg
    elif name == "resnet101":
        from torchvision.models import resnet101 as Alg
    elif name == "resnet152":
        from torchvision.models import resnet152 as Alg
    elif name == "resnet18":
        from torchvision.models import resnet18 as Alg
    elif name == "resnet34":
        from torchvision.models import resnet34 as Alg
    elif name == "resnet50":
        from torchvision.models import resnet50 as Alg
    elif name == "resnext101_32x8d":
        from torchvision.models import resnext101_32x8d as Alg
    elif name == "resnext101_64x4d":
        from torchvision.models import resnext101_64x4d as Alg
    elif name == "resnext50_32x4d":
        from torchvision.models import resnext50_32x4d as Alg
    elif name == "shufflenet_v2_x0_5":
        from torchvision.models import shufflenet_v2_x0_5 as Alg
    elif name == "shufflenet_v2_x1_0":
        from torchvision.models import shufflenet_v2_x1_0 as Alg
    elif name == "shufflenet_v2_x1_5":
        from torchvision.models import shufflenet_v2_x1_5 as Alg
    elif name == "shufflenet_v2_x2_0":
        from torchvision.models import shufflenet_v2_x2_0 as Alg
    elif name == "squeezenet1_0":
        from torchvision.models import squeezenet1_0 as Alg
    elif name == "squeezenet1_1":
        from torchvision.models import squeezenet1_1 as Alg
    elif name == "swin_b":
        from torchvision.models import swin_b as Alg
    elif name == "swin_s":
        from torchvision.models import swin_s as Alg
    elif name == "swin_t":
        from torchvision.models import swin_t as Alg
    elif name == "vgg11":
        from torchvision.models import vgg11 as Alg
    elif name == "vgg11_bn":
        from torchvision.models import vgg11_bn as Alg
    elif name == "vgg13":
        from torchvision.models import vgg13 as Alg
    elif name == "vgg13_bn":
        from torchvision.models import vgg13_bn as Alg
    elif name == "vgg16":
        from torchvision.models import vgg16 as Alg
    elif name == "vgg16_bn":
        from torchvision.models import vgg16_bn as Alg
    elif name == "vgg19":
        from torchvision.models import vgg19 as Alg
    elif name == "vgg19_bn":
        from torchvision.models import vgg19_bn as Alg
    elif name == "vit_b_16":
        from torchvision.models import vit_b_16 as Alg
    elif name == "vit_b_32":
        from torchvision.models import vit_b_32 as Alg
    elif name == "vit_h_14":
        from torchvision.models import vit_h_14 as Alg
    elif name == "vit_l_16":
        from torchvision.models import vit_l_16 as Alg
    elif name == "vit_l_32":
        from torchvision.models import vit_l_32 as Alg
    elif name == "wide_resnet101_2":
        from torchvision.models import wide_resnet101_2 as Alg
    elif name == "wide_resnet50_2":
        from torchvision.models import wide_resnet50_2 as Alg
    return Alg
