def get_model(model, input_channels=3, pretrained=False):
if model == 'resnet50_md':
out_model = Resnet50_md(input_channels)
elif model == 'resnet18_md':
out_model = Resnet18_md(input_channels)
else:
out_model = ResnetModel(input_channels, encoder=model, pretrained=pretrained)
return out_model
def get_model(model, input_channels=3, pretrained=True):
if model == 'resnet50_md':
out_model = Resnet50_md(input_channels)
elif model == 'resnet18_md':
out_model = Resnet18_md(input_channels)
elif model == 'AlbUNet':
out_model = UNet()
elif model == 'nested':
out_model = NestedUNet()
elif model == 'ldid':
out_model = Model(ResidualBlock, UpProj_Block, 8)
else:
out_model = ResnetModel(input_channels,
encoder=model,
pretrained=pretrained)
return out_model