def __init__(self, pretrained, num_classes):
super(SegNet, self).__init__()
vgg = models.vgg19_bn()
if pretrained:
vgg.load_state_dict(torch.load(pretrained_vgg19_bn))
features = list(vgg.features.children())
self.enc1 = nn.Sequential(*features[0:7])
self.enc2 = nn.Sequential(*features[7:14])
self.enc3 = nn.Sequential(*features[14:27])
self.enc4 = nn.Sequential(*features[27:40])
self.enc5 = nn.Sequential(*features[40:])
# for m in self.modules():
# if isinstance(m, nn.Conv2d):
# m.requires_grad = False
self.dec5 = nn.Sequential(
*([nn.ConvTranspose2d(512, 512, kernel_size=2, stride=2)] +
[nn.Conv2d(512, 512, kernel_size=3, padding=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)] * 4)
)
self.dec4 = _DecoderBlock(1024, 256, 4)
self.dec3 = _DecoderBlock(512, 128, 4)
self.dec2 = _DecoderBlock(256, 64, 2)
self.dec1 = _DecoderBlock(128, num_classes, 2)
initialize_weights(self.dec5, self.dec4, self.dec3, self.dec2, self.dec1)
def __init__(self, num_classes, pretrained=True):
super(FCN32DenseNet, self).__init__()
dense = models.densenet201()
if pretrained:
dense.load_state_dict(torch.load(dense201_path))
self.features5 = dense.features
self.fconv5 = nn.Sequential(
nn.ReLU(inplace=True), nn.Conv2d(1920, num_classes, kernel_size=7))
initialize_weights(self.fconv5)
def __init__(self, num_classes, pretrained=True):
super(FCN32ResNet, self).__init__()
res = models.resnet152()
if pretrained:
res.load_state_dict(torch.load(res152_path))
self.features5 = nn.Sequential(res.conv1, res.bn1, res.relu,
res.maxpool, res.layer1, res.layer2,
res.layer3, res.layer4)
self.fconv5 = nn.Conv2d(2048, num_classes, kernel_size=7)
initialize_weights(self.fconv5)
def __init__(self, pretrained, num_classes):
super(FCN16ResNet, self).__init__()
res = models.resnet152()
if pretrained:
res.load_state_dict(torch.load(pretrained_res152))
self.features4 = nn.Sequential(
res.conv1, res.bn1, res.relu, res.maxpool, res.layer1, res.layer2, res.layer3
)
self.features5 = res.layer4
self.fconv4 = nn.Conv2d(1024, num_classes, kernel_size=1)
self.fconv5 = nn.Conv2d(2048, num_classes, kernel_size=7)
initialize_weights(self.fconv4, self.fconv5)
def __init__(self, num_classes, pretrained=True):
super(FCN32VGG, self).__init__()
vgg = models.vgg19_bn()
if pretrained:
vgg.load_state_dict(torch.load(vgg19_bn_path))
self.features5 = vgg.features
self.fconv5 = nn.Sequential(
nn.Conv2d(512, 4096, kernel_size=7), nn.ReLU(inplace=True),
nn.Dropout(), nn.Conv2d(4096, 4096, kernel_size=1),
nn.ReLU(inplace=True), nn.Dropout(),
nn.Conv2d(4096, num_classes, kernel_size=1))
initialize_weights(self.fconv5)
def __init__(self, num_classes, pretrained=True):
super(FCN16DenseNet, self).__init__()
dense = models.densenet201()
if pretrained:
dense.load_state_dict(torch.load(dense201_path))
features = list(dense.features.children())
self.features4 = nn.Sequential(*features[:10])
self.features5 = nn.Sequential(*features[10:])
self.fconv4 = nn.Sequential(nn.BatchNorm2d(896), nn.ReLU(inplace=True),
nn.Conv2d(896, num_classes, kernel_size=1))
self.fconv5 = nn.Sequential(
nn.ReLU(inplace=True), nn.Conv2d(1920, num_classes, kernel_size=7))
initialize_weights(self.fconv4, self.fconv5)
def __init__(self, num_classes, input_size, pretrained=True, use_aux=True):
super(PSPNetDeform, self).__init__()
self.input_size = input_size
self.use_aux = use_aux
resnet = models.resnet152()
if pretrained:
resnet.load_state_dict(torch.load(res152_path))
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu,
resnet.maxpool)
self.layer1 = resnet.layer1
self.layer2 = resnet.layer2
self.layer3 = resnet.layer3
self.layer4 = resnet.layer4
for n, m in self.layer3.named_modules():
if 'conv2' in n:
m.padding = (1, 1)
m.stride = (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for n, m in self.layer4.named_modules():
if 'conv2' in n:
m.padding = (1, 1)
m.stride = (1, 1)
elif 'downsample.0' in n:
m.stride = (1, 1)
for idx in range(len(self.layer3)):
self.layer3[idx].conv2 = Conv2dDeformable(self.layer3[idx].conv2)
for idx in range(len(self.layer4)):
self.layer4[idx].conv2 = Conv2dDeformable(self.layer4[idx].conv2)
self.ppm = PyramidPoolingModule((int(math.ceil(
input_size[0] / 8.0)), int(math.ceil(input_size[1] / 8.0))), 2048,
512, (1, 2, 3, 6))
self.final = nn.Sequential(
nn.Conv2d(4096, 512, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(512, momentum=.95), nn.ReLU(), nn.Dropout(0.1),
nn.Conv2d(512, num_classes, kernel_size=1))
if use_aux:
self.aux_logits = nn.Sequential(
PyramidPoolingModule((int(math.ceil(input_size[0] / 8.0)),
int(math.ceil(input_size[1] / 8.0))),
1024, 256, (1, 2, 3, 6)),
nn.Conv2d(2048, 256, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(256, momentum=.95), nn.ReLU(), nn.Dropout(0.1),
nn.Conv2d(256, num_classes, kernel_size=1))
initialize_weights(self.ppm, self.final)
def __init__(self, pretrained, num_classes):
super(FCN8VGG, self).__init__()
vgg = models.vgg19_bn()
if pretrained:
vgg.load_state_dict(torch.load(pretrained_vgg19_bn))
features = list(vgg.features.children())
self.features3 = nn.Sequential(*features[0:27])
self.features4 = nn.Sequential(*features[27:40])
self.features5 = nn.Sequential(*features[40:])
self.fconv3 = nn.Conv2d(256, num_classes, kernel_size=1)
self.fconv4 = nn.Conv2d(512, num_classes, kernel_size=1)
self.fconv5 = nn.Sequential(
nn.Conv2d(512, 4096, kernel_size=7), nn.ReLU(inplace=True),
nn.Dropout(), nn.Conv2d(4096, 4096, kernel_size=1),
nn.ReLU(inplace=True), nn.Dropout(),
nn.Conv2d(4096, num_classes, kernel_size=1))
initialize_weights(self.fconv3, self.fconv4, self.fconv5)
def __init__(self, num_classes):
super(UNet, self).__init__()
self.enc1 = _EncoderBlock(3, 64)
self.enc2 = _EncoderBlock(64, 128)
self.enc3 = _EncoderBlock(128, 256)
self.enc4 = _EncoderBlock(256, 512, dropout=True)
self.center = _DecoderBlock(512, 1024, 512)
self.dec4 = _DecoderBlock(1024, 512, 256)
self.dec3 = _DecoderBlock(512, 256, 128)
self.dec2 = _DecoderBlock(256, 128, 64)
self.dec1 = nn.Sequential(
nn.Conv2d(128, 64, kernel_size=3),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, kernel_size=3),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
self.final = nn.Conv2d(64, num_classes, kernel_size=1)
initialize_weights(self)
def __init__(self, num_classes, input_size, pretrained=True):
super(GCN, self).__init__()
self.input_size = input_size
resnet = models.resnet152()
if pretrained:
resnet.load_state_dict(torch.load(res152_path))
self.layer0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu)
self.layer1 = nn.Sequential(resnet.maxpool, resnet.layer1)
self.layer2 = resnet.layer2
self.layer3 = resnet.layer3
self.layer4 = resnet.layer4
# use large kernel with odd size
gcm_ks = [[input_size[0] / 32 - 1, input_size[1] / 32 - 1],
[input_size[0] / 16 - 1, input_size[1] / 16 - 1],
[input_size[0] / 8 - 1, input_size[1] / 8 - 1],
[input_size[0] / 4 - 1, input_size[1] / 4 - 1]]
gcm_ks = [(ks[0] - 1 if ks[0] % 2 == 0 else ks[0],
ks[1] - 1 if ks[1] % 2 == 0 else ks[1]) for ks in gcm_ks]
self.gcm1 = GlobalConvModule(2048, num_classes, gcm_ks[0])
self.gcm2 = GlobalConvModule(1024, num_classes, gcm_ks[1])
self.gcm3 = GlobalConvModule(512, num_classes, gcm_ks[2])
self.gcm4 = GlobalConvModule(256, num_classes, gcm_ks[3])
self.brm1 = BoundaryRefineModule(num_classes)
self.brm2 = BoundaryRefineModule(num_classes)
self.brm3 = BoundaryRefineModule(num_classes)
self.brm4 = BoundaryRefineModule(num_classes)
self.brm5 = BoundaryRefineModule(num_classes)
self.brm6 = BoundaryRefineModule(num_classes)
self.brm7 = BoundaryRefineModule(num_classes)
self.brm8 = BoundaryRefineModule(num_classes)
self.brm9 = BoundaryRefineModule(num_classes)
initialize_weights(self.gcm1, self.gcm2, self.gcm3, self.gcm4,
self.brm1, self.brm2, self.brm3, self.brm4,
self.brm5, self.brm6, self.brm7, self.brm8,
self.brm9)
