def hrnet48_unet32(input_channels=3,
num_classes=1,
dropout=0.0,
pretrained=True):
encoder = E.HRNetV2Encoder48(pretrained=pretrained)
if input_channels != 3:
encoder.change_input_channels(input_channels)
return UnetSegmentationModel(encoder,
num_classes=num_classes,
unet_channels=[32, 64, 128, 256],
dropout=dropout)
def hrnet48_unet64(input_channels=3, num_classes=1, dropout=0.0, pretrained=True):
encoder = E.HRNetV2Encoder48(pretrained=pretrained, layers=[1, 2, 3, 4])
if input_channels != 3:
encoder.change_input_channels(input_channels)
return UnetV3SegmentationModel(
encoder,
num_classes=num_classes,
unet_channels=[128, 128, 256],
last_upsample_filters=64,
dropout=dropout,
abn_block=partial(ABN, activation=ACT_RELU),
)
def hrnet48(num_classes=1, dropout=0.0, pretrained=True):
encoder = E.HRNetV2Encoder48(pretrained=pretrained)
return HRNetSegmentationModel(encoder, num_classes=num_classes, dropout=dropout)