def __init__(self, model_type, num_class, device, pre_model):
self.model_type = model_type
self.num_class = num_class
if self.model_type == "VGG16_Finetuning":
self.model = VGG16_Finetuning.VGG16_Finetuning(
self.num_class).to(device)
if self.model_type == 'vgg16':
self.model = vgg16.VGG16(self.num_class).to(device)
self.model = self.load_from_pretrain(self.model, pre_model)
if self.model_type == 'resnet':
self.model = ResNet.ResNet50(self.num_class).to(device)
def __call__(self, x, seon):
return {
1: Vgg.VGG11(seon),
2: Vgg.VGG13(seon),
3: Vgg.VGG16(seon),
4: Vgg.VGG19(seon),
5: ResNet.ResNet18(seon),
6: ResNet.ResNet34(seon),
7: ResNet.ResNet50(seon),
8: ResNet.ResNet101(seon),
9: ResNet.ResNet152(seon),
10: DenseNet.DenseNet121(seon),
11: DenseNet.DenseNet169(seon),
12: DenseNet.DenseNet201(seon),
13: DenseNet.DenseNet161(seon)
}[x]
def BackBone_Unet(backbone_name):
up_parm_dict = {
'resnet18': [512, 256, 128, 64, 64, 64, 64, 3],
'resnet34': [512, 256, 128, 64, 64, 64, 64, 3],
'resnet50': [2048, 1024, 512, 256, 128, 64, 64, 3],
'resnet101': [2048, 1024, 512, 256, 128, 64, 64, 3],
'resnet152': [2048, 1024, 512, 256, 128, 64, 64, 3],
'densenet121': [1024, 1024, 512, 256, 128, 64, 64, 3],
'densenet161': [2204, 2104, 752, 352, 128, 64, 64, 3],
'densenet201': [1920, 1792, 512, 256, 128, 64, 64, 3],
'densenet169': [1664, 1280, 512, 256, 128, 64, 64, 3],
'efficientnet-b0': [1280, 112, 40, 24, 16, 16, 64, 3],
'efficientnet-b1': [1280, 112, 40, 24, 16, 16, 64, 3],
'efficientnet-b2': [1280, 120, 48, 24, 16, 16, 64, 3],
'efficientnet-b3': [1280, 136, 48, 32, 24, 24, 64, 3],
'efficientnet-b4': [1280, 160, 56, 32, 24, 24, 64, 3],
'efficientnet-b5': [1280, 176, 64, 40, 24, 24, 64, 3],
'efficientnet-b6': [1280, 200, 72, 40, 32, 32, 64, 3],
'efficientnet-b7': [1280, 224, 80, 48, 32, 32, 64, 3]
}
efficient_param = {
# 'efficientnet type': (width_coef, depth_coef, resolution, dropout_rate)
'efficientnet-b0': (1.0, 1.0, 224, 0.2),
'efficientnet-b1': (1.0, 1.1, 224, 0.2),
'efficientnet-b2': (1.1, 1.2, 224, 0.3),
'efficientnet-b3': (1.2, 1.4, 224, 0.3),
'efficientnet-b4': (1.4, 1.8, 224, 0.4),
'efficientnet-b5': (1.6, 2.2, 224, 0.4),
'efficientnet-b6': (1.8, 2.6, 224, 0.5),
'efficientnet-b7': (2.0, 3.1, 224, 0.5)
}
if backbone_name[0] == 'r':
if backbone_name[-2:] == '18':
model = ResNet.ResNet18()
if backbone_name[-2:] == '34':
model = ResNet.ResNet34()
if backbone_name[-2:] == '50':
model = ResNet.ResNet50()
if backbone_name[-2:] == '01':
model = ResNet.ResNet101()
if backbone_name[-2:] == '52':
model = ResNet.ResNet152()
net = Res_Unet(model=model, up_parm=up_parm_dict[backbone_name])
elif backbone_name[0] == 'd':
if backbone_name[-2:] == '21':
model = DenseNet.DenseNet121(seon=False)
if backbone_name[-2:] == '61':
model = DenseNet.DenseNet161(seon=False)
if backbone_name[-2:] == '01':
model = DenseNet.DenseNet201(seon=False)
if backbone_name[-2:] == '69':
model = DenseNet.DenseNet169(seon=False)
net = Dense_Unet(model=model, up_parm=up_parm_dict[backbone_name])
elif backbone_name[0] == 'e':
param = efficient_param[backbone_name]
model = EfficientNet.EfficientNet(param)
net = Efficient_Unet(model=model, up_parm=up_parm_dict[backbone_name])
return net
