# Source: https://github.com/zijundeng/pytorch-semantic-segmentation/tree/master/models (MIT)
"""
Implementation of `Segnet: A deep convolutional encoder-decoder architecture for image segmentation <https://arxiv.org/pdf/1511.00561>`_
"""
import torch
from torch import nn
from torchvision import models
from .fcn_utils import initialize_weights
from .config import vgg19_bn_path
__all__ = ['SegNet']
class _DecoderBlock(nn.Module):
def __init__(self, in_channels, out_channels, num_conv_layers):
super(_DecoderBlock, self).__init__()
middle_channels = in_channels / 2
layers = [
nn.ConvTranspose2d(in_channels, in_channels, kernel_size=2, stride=2),
nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(middle_channels),
nn.ReLU(inplace=True)
]
layers += [
nn.Conv2d(middle_channels, middle_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(middle_channels),
nn.ReLU(inplace=True),
] * (num_conv_layers - 2)
layers += [
nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True),
]
self.decode = nn.Sequential(*layers)
def forward(self, x):
return self.decode(x)
[docs]class SegNet(nn.Module):
def __init__(self, num_classes, pretrained=True, **kwargs):
super(SegNet, self).__init__()
vgg = models.vgg19_bn()
if pretrained:
vgg.load_state_dict(torch.load(vgg19_bn_path))
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:])
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 forward(self, x):
enc1 = self.enc1(x)
enc2 = self.enc2(enc1)
enc3 = self.enc3(enc2)
enc4 = self.enc4(enc3)
enc5 = self.enc5(enc4)
dec5 = self.dec5(enc5)
dec4 = self.dec4(torch.cat([enc4, dec5], 1))
dec3 = self.dec3(torch.cat([enc3, dec4], 1))
dec2 = self.dec2(torch.cat([enc2, dec3], 1))
dec1 = self.dec1(torch.cat([enc1, dec2], 1))
return dec1