def __init__(self, in_channels, num_categories=2, deep_supervision=False, pretrained=True):
super(ResNet34_UNet, self).__init__()
self.in_channels = in_channels
self.num_categories = num_categories
self.deep_supervision = deep_supervision
self.pretrained = pretrained
self.filter_sizes = None
self.main_model_name = 'resnet34_unet'
# Encoder:
self.enc1 = nn.ModuleList([resnet34(pretrained=self.pretrained).conv1, resnet34(pretrained=self.pretrained).bn1,
resnet34(pretrained=self.pretrained).relu,
resnet34(pretrained=self.pretrained).maxpool])
self.enc2 = nn.ModuleList([resnet34(pretrained=self.pretrained).layer1])
self.enc3 = nn.ModuleList([resnet34(pretrained=self.pretrained).layer2])
self.enc4 = nn.ModuleList([resnet34(pretrained=self.pretrained).layer3])
self.enc5 = nn.ModuleList([resnet34(pretrained=self.pretrained).layer4])
# Decoder:
self.dec4_UpConcat2d = UpConcat2d(512, 256)
self.dec4_Conv2dBlock = Conv2dBlock(512, 256, True)
self.dec3_UpConcat2d = UpConcat2d(256, 128)
self.dec3_Conv2dBlock = Conv2dBlock(256, 128, True)
self.dec2_UpConcat2d = UpConcat2d(128, 64)
self.dec2_Conv2dBlock = Conv2dBlock(128, 64, True)
self.dec1_Concat2d = Concat2d()
self.dec1_Conv2dBlock = Conv2dBlock(128, 64, True)
if self.deep_supervision:
aggr_channels = 512
self.dec1_out = nn.Conv2d(in_channels=aggr_channels, out_channels=self.num_categories, kernel_size=1,
stride=1, padding=0)
self.deep_sup_module = DeepSupervisionModule()
else:
self.dec1_out = nn.Conv2d(in_channels=64, out_channels=self.num_categories, kernel_size=1,
stride=1, padding=0)
def __init__(self, in_channels, num_categories=2, deep_supervision=False, pretrained=True):
super(MobileNetV2_UNet, self).__init__()
self.in_channels = in_channels
self.num_categories = num_categories
self.deep_supervision = deep_supervision
self.pretrained = pretrained
self.filter_sizes = None
self.main_model_name = 'mobilenetv2_unet'
# Encoder:
self.enc1 = nn.ModuleList(list(mobilenet_v2(pretrained=self.pretrained).features)[:4])
self.enc2 = nn.ModuleList(list(mobilenet_v2(pretrained=self.pretrained).features)[4:11])
self.enc3 = nn.ModuleList(list(mobilenet_v2(pretrained=self.pretrained).features)[11:14])
self.enc4 = nn.ModuleList(list(mobilenet_v2(pretrained=self.pretrained).features)[14:17])
self.enc5 = nn.ModuleList([list(mobilenet_v2(pretrained=self.pretrained).features)[17]])
# Decoder:
self.dec4_Concat2d = Concat2d()
self.dec4_Conv2dBlock = Conv2dBlock(480, 160, True)
self.dec3_UpConcat2d = UpConcat2d(160, 96)
self.dec3_Conv2dBlock = Conv2dBlock(192, 96, True)
self.dec2_Concat2d = Concat2d()
self.dec2_Conv2dBlock = Conv2dBlock(160, 64, True)
self.dec1_UpConcat2d = UpConcat2d(64, 24, scale_factor=4)
self.dec1_Conv2dBlock = Conv2dBlock(48, 24, True)
if self.deep_supervision:
aggr_channels = 344
self.dec1_out = nn.Conv2d(in_channels=aggr_channels, out_channels=self.num_categories, kernel_size=1,
stride=1, padding=0)
self.deep_sup_module = DeepSupervisionModule(up_sampling_factors=(2, 2, 4))
else:
self.dec1_out = nn.Conv2d(in_channels=24, out_channels=self.num_categories, kernel_size=1,
stride=1, padding=0)
def __init__(self,
in_channels,
num_categories=2,
filter_sizes=(64, 128, 256, 512, 1024),
deep_supervision=False):
super(UNet, self).__init__()
self.in_channels = in_channels
self.num_categories = num_categories
self.filter_sizes = filter_sizes
self.deep_supervision = deep_supervision
self.main_model_name = 'unet'
# Encoder:
self.enc1_Conv2dBlock = Conv2dBlock(self.in_channels,
self.filter_sizes[0], True)
self.enc2_Conv2dBlock = Conv2dBlock(self.filter_sizes[0],
self.filter_sizes[1], True)
self.enc3_Conv2dBlock = Conv2dBlock(self.filter_sizes[1],
self.filter_sizes[2], True)
self.enc4_Conv2dBlock = Conv2dBlock(self.filter_sizes[2],
self.filter_sizes[3], True)
self.enc5_Conv2dBlock = Conv2dBlock(self.filter_sizes[3],
self.filter_sizes[4], True)
self.pool = nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2), padding=0)
# Decoder:
self.dec4_UpConcat2d = UpConcat2d(self.filter_sizes[4],
self.filter_sizes[3])
self.dec4_Conv2dBlock = Conv2dBlock(self.filter_sizes[4],
self.filter_sizes[3], True)
self.dec3_UpConcat2d = UpConcat2d(self.filter_sizes[3],
self.filter_sizes[2])
self.dec3_Conv2dBlock = Conv2dBlock(self.filter_sizes[3],
self.filter_sizes[2], True)
self.dec2_UpConcat2d = UpConcat2d(self.filter_sizes[2],
self.filter_sizes[1])
self.dec2_Conv2dBlock = Conv2dBlock(self.filter_sizes[2],
self.filter_sizes[1], True)
self.dec1_UpConcat2d = UpConcat2d(self.filter_sizes[1],
self.filter_sizes[0])
self.dec1_Conv2dBlock = Conv2dBlock(self.filter_sizes[1],
self.filter_sizes[0], True)
if self.deep_supervision:
aggr_channels = self.filter_sizes[0] + self.filter_sizes[
1] + self.filter_sizes[2] + self.filter_sizes[3]
self.dec1_out = nn.Conv2d(in_channels=aggr_channels,
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
self.deep_sup_module = DeepSupervisionModule()
else:
self.dec1_out = nn.Conv2d(in_channels=self.filter_sizes[0],
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
def __init__(self,
in_channels,
num_categories=2,
deep_supervision=False,
pretrained=True):
super(vgg11_UNet, self).__init__()
self.in_channels = in_channels
self.num_categories = num_categories
self.deep_supervision = deep_supervision
self.pretrained = pretrained
self.filter_sizes = None
self.main_model_name = 'vgg11_unet'
# Encoder:
self.enc1 = nn.ModuleList(
list(vgg11(pretrained=self.pretrained).features)[:2])
self.enc2 = nn.ModuleList(
list(vgg11(pretrained=self.pretrained).features)[2:5])
self.enc3 = nn.ModuleList(
list(vgg11(pretrained=self.pretrained).features)[5:10])
self.enc4 = nn.ModuleList(
list(vgg11(pretrained=self.pretrained).features)[10:15])
self.enc5 = nn.ModuleList(
list(vgg11(pretrained=self.pretrained).features)[15:20])
# Decoder:
self.dec4_UpConcat2d = UpConcat2d(512, 256)
self.dec4_Conv2dBlock = Conv2dBlock(768, 256, True)
self.dec3_UpConcat2d = UpConcat2d(256, 128)
self.dec3_Conv2dBlock = Conv2dBlock(384, 256, True)
self.dec2_UpConcat2d = UpConcat2d(256, 128)
self.dec2_Conv2dBlock = Conv2dBlock(256, 128, True)
self.dec1_UpConcat2d = UpConcat2d(128, 64)
self.dec1_Conv2dBlock = Conv2dBlock(128, 64, True)
if self.deep_supervision:
aggr_channels = 704
self.dec1_out = nn.Conv2d(in_channels=aggr_channels,
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
self.deep_sup_module = DeepSupervisionModule()
else:
self.dec1_out = nn.Conv2d(in_channels=64,
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
def __init__(self,
in_channels,
num_categories=2,
deep_supervision=False,
pretrained=True):
super(DenseNet121_UNet, self).__init__()
self.in_channels = in_channels
self.num_categories = num_categories
self.deep_supervision = deep_supervision
self.pretrained = pretrained
self.filter_sizes = None
self.main_model_name = 'densenet121_unet'
# Encoder:
self.enc1 = nn.ModuleList(
list(densenet121(pretrained=self.pretrained).features)[:3])
self.enc2 = nn.ModuleList(
list(densenet121(pretrained=self.pretrained).features)[3:5])
self.enc3 = nn.ModuleList(
list(densenet121(pretrained=self.pretrained).features)[5:7])
self.enc4 = nn.ModuleList(
list(densenet121(pretrained=self.pretrained).features)[7:9])
self.enc5 = nn.ModuleList(
list(densenet121(pretrained=self.pretrained).features)[9:11])
# Decoder:
self.dec4_UpConcat2d = UpConcat2d(1024, 1024)
self.dec4_Conv2dBlock = Conv2dBlock(2048, 1024, True)
self.dec3_UpConcat2d = UpConcat2d(1024, 512)
self.dec3_Conv2dBlock = Conv2dBlock(1024, 512, True)
self.dec2_UpConcat2d = UpConcat2d(512, 256)
self.dec2_Conv2dBlock = Conv2dBlock(512, 256, True)
self.dec1_UpConcat2d = UpConcat2d(256, 64)
self.dec1_Conv2dBlock = Conv2dBlock(128, 64, True)
if self.deep_supervision:
aggr_channels = 1856
self.dec1_out = nn.Conv2d(in_channels=aggr_channels,
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
self.deep_sup_module = DeepSupervisionModule()
else:
self.dec1_out = nn.Conv2d(in_channels=64,
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
def __init__(self,
in_channels,
num_categories=2,
filter_sizes=(64, 128, 256, 512, 1024),
deep_supervision=True):
super(ResidualAttentionUNet, self).__init__()
self.in_channels = in_channels
self.num_categories = num_categories
self.filter_sizes = filter_sizes
self.deep_supervision = deep_supervision
self.main_model_name = 'residualattention_unet'
# Encoder:
self.enc1_ResNeXt_Block = ResNeXt_Block(self.in_channels,
self.filter_sizes[0])
self.enc2_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[0],
self.filter_sizes[1])
self.enc3_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[1],
self.filter_sizes[2])
self.enc4_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[2],
self.filter_sizes[3])
self.enc5_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[3],
self.filter_sizes[4])
self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
# Decoder:
self.dec4_att_gate = AttentionGate(self.filter_sizes[3],
self.filter_sizes[4])
self.dec4_UpConcat2d = UpConcat2d(self.filter_sizes[4],
self.filter_sizes[3])
self.dec4_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[4],
self.filter_sizes[3])
self.dec3_att_gate = AttentionGate(self.filter_sizes[2],
self.filter_sizes[3])
self.dec3_UpConcat2d = UpConcat2d(self.filter_sizes[3],
self.filter_sizes[2])
self.dec3_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[3],
self.filter_sizes[2])
self.dec2_att_gate = AttentionGate(self.filter_sizes[1],
self.filter_sizes[2])
self.dec2_UpConcat2d = UpConcat2d(self.filter_sizes[2],
self.filter_sizes[1])
self.dec2_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[2],
self.filter_sizes[1])
self.dec1_att_gate = AttentionGate(self.filter_sizes[0],
self.filter_sizes[1])
self.dec1_UpConcat2d = UpConcat2d(self.filter_sizes[1],
self.filter_sizes[0])
self.dec1_ResNeXt_Block = ResNeXt_Block(self.filter_sizes[1],
self.filter_sizes[0])
if self.deep_supervision:
aggr_channels = self.filter_sizes[0] + self.filter_sizes[
1] + self.filter_sizes[2] + self.filter_sizes[3]
self.dec1_out = nn.Conv2d(in_channels=aggr_channels,
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)
self.deep_sup_module = DeepSupervisionModule()
else:
self.dec1_out = nn.Conv2d(in_channels=self.filter_sizes[0],
out_channels=self.num_categories,
kernel_size=1,
stride=1,
padding=0)