def create_model(device, type ='vgg16'):
assert type == 'vgg16' or type == 'resnet101'
if type == 'vgg16':
model = UNet16(pretrained=True)
elif type == 'resnet101':
model = UNetResNet(pretrained=True, encoder_depth=101, num_classes=1)
else:
assert False
model.eval()
return model.to(device)
def load_unet_resnet_34(model_path):
model = UNetResNet(pretrained=True, encoder_depth=34, num_classes=1)
checkpoint = torch.load(model_path)
if 'model' in checkpoint:
model.load_state_dict(checkpoint['model'])
elif 'state_dict' in checkpoint:
model.load_state_dict(checkpoint['check_point'])
else:
raise Exception('undefind model format')
model.cuda()
model.eval()
return model
def create_model(device, type='vgg16'):
assert type == 'vgg16' or type == 'resnet101'
if type == 'vgg16':
print('create vgg16 model')
model = UNet16(pretrained=True)
elif type == 'resnet101':
encoder_depth = 101
num_classes = 1
print('create resnet101 model')
model = UNetResNet(encoder_depth=encoder_depth,
num_classes=num_classes,
pretrained=True)
else:
assert False
model.eval()
return model.to(device)
def create_model(device, type ='vgg16'):
if type == 'vgg16':
print('create vgg16 model')
model = UNet16(pretrained=True)
elif type == 'vgg16_bn':
print('create vgg16_bn model')
model = UNet16_bn(pretrained=True)
elif type == 'vgg16_bn_do':
print('create vgg16_bn_do model')
model = UNet16_bn_do(pretrained=True)
elif type == 'vgg16_fullbn_do':
print('create vgg16_fullbn_do model')
model = UNet16_fullbn_do(pretrained=True)
elif type == 'resnet101':
encoder_depth = 101
num_classes = 1
print('create resnet101 model')
model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True)
elif type == 'resnet34':
encoder_depth = 34
num_classes = 1
print('create resnet34 model')
model = UNetResNet(encoder_depth=encoder_depth, num_classes=num_classes, pretrained=True)
elif type == 'unet+++':
model = UNet_3Plus()
print('create U-Net+++ model')
elif type =='unet++':
model = UNet_2Plus()
print('create U-Net++ model')
else:
assert False
model.eval()
return model.to(device)