Update unet.py

unet.py

@@ -1,65 +1,67 @@
-import torch
 import torch.nn as nn
+from model_utils import *
 
-class UNet(nn.Module):
-    def __init__(self, n_channels=3, n_classes=19):
-        super(UNet, self).__init__()
-        
-        def CBR(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
-            return nn.Sequential(
-                nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias),
-                nn.BatchNorm2d(out_channels),
-                nn.ReLU(inplace=True)
-            )
-        
-        # Encoder
-        self.enc1 = nn.Sequential(CBR(n_channels, 64), CBR(64, 64))
-        self.enc2 = nn.Sequential(nn.MaxPool2d(2), CBR(64, 128), CBR(128, 128))
-        self.enc3 = nn.Sequential(nn.MaxPool2d(2), CBR(128, 256), CBR(256, 256))
-        self.enc4 = nn.Sequential(nn.MaxPool2d(2), CBR(256, 512), CBR(512, 512))
-        self.enc5 = nn.Sequential(nn.MaxPool2d(2), CBR(512, 1024), CBR(1024, 1024))
-        
-        # Decoder
-        self.dec4 = nn.Sequential(CBR(1024+512, 512), CBR(512, 512))
-        self.dec3 = nn.Sequential(CBR(512+256, 256), CBR(256, 256))
-        self.dec2 = nn.Sequential(CBR(256+128, 128), CBR(128, 128))
-        self.dec1 = nn.Sequential(CBR(128+64, 64), CBR(64, 64))
-        
-        # Upsampling
-        self.up4 = nn.ConvTranspose2d(1024, 512, 2, 2)
-        self.up3 = nn.ConvTranspose2d(512, 256, 2, 2)
-        self.up2 = nn.ConvTranspose2d(256, 128, 2, 2)
-        self.up1 = nn.ConvTranspose2d(128, 64, 2, 2)
-        
-        # Final layer
-        self.final = nn.Conv2d(64, n_classes, 1)
-        
-    def forward(self, x):
-        # Encoder
-        e1 = self.enc1(x)
-        e2 = self.enc2(e1)
-        e3 = self.enc3(e2)
-        e4 = self.enc4(e3)
-        e5 = self.enc5(e4)
-        
-        # Decoder with skip connections
-        d4 = self.up4(e5)
-        d4 = torch.cat([d4, e4], dim=1)
-        d4 = self.dec4(d4)
-        
-        d3 = self.up3(d4)
-        d3 = torch.cat([d3, e3], dim=1)
-        d3 = self.dec3(d3)
-        
-        d2 = self.up2(d3)
-        d2 = torch.cat([d2, e2], dim=1)
-        d2 = self.dec2(d2)
-        
-        d1 = self.up1(d2)
-        d1 = torch.cat([d1, e1], dim=1)
-        d1 = self.dec1(d1)
-        
-        return self.final(d1)
-
-def unet(**kwargs):
-    return UNet(**kwargs)
+class unet(nn.Module):
+    def __init__(
+        self,
+        feature_scale=4,
+        n_classes=19,
+        is_deconv=True,
+        in_channels=3,
+        is_batchnorm=True,
+    ):
+        super(unet, self).__init__()
+        self.is_deconv = is_deconv
+        self.in_channels = in_channels
+        self.is_batchnorm = is_batchnorm
+        self.feature_scale = feature_scale
+
+        filters = [64, 128, 256, 512, 1024]
+        filters = [int(x / self.feature_scale) for x in filters]
+
+        # downsampling
+        self.conv1 = unetConv2(self.in_channels, filters[0], self.is_batchnorm)
+        self.maxpool1 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv2 = unetConv2(filters[0], filters[1], self.is_batchnorm)
+        self.maxpool2 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv3 = unetConv2(filters[1], filters[2], self.is_batchnorm)
+        self.maxpool3 = nn.MaxPool2d(kernel_size=2)
+
+        self.conv4 = unetConv2(filters[2], filters[3], self.is_batchnorm)
+        self.maxpool4 = nn.MaxPool2d(kernel_size=2)
+
+        self.center = unetConv2(filters[3], filters[4], self.is_batchnorm)
+
+        # upsampling
+        self.up_concat4 = unetUp(filters[4], filters[3], self.is_deconv, self.is_batchnorm)
+        self.up_concat3 = unetUp(filters[3], filters[2], self.is_deconv, self.is_batchnorm)
+        self.up_concat2 = unetUp(filters[2], filters[1], self.is_deconv, self.is_batchnorm)
+        self.up_concat1 = unetUp(filters[1], filters[0], self.is_deconv, self.is_batchnorm)
+
+        # final conv (without any concat)
+        self.final = nn.Conv2d(filters[0], n_classes, 1)
+
+    def forward(self, inputs):
+        conv1 = self.conv1(inputs)
+        maxpool1 = self.maxpool1(conv1)
+
+        conv2 = self.conv2(maxpool1)
+        maxpool2 = self.maxpool2(conv2)
+
+        conv3 = self.conv3(maxpool2)
+        maxpool3 = self.maxpool3(conv3)
+
+        conv4 = self.conv4(maxpool3)
+        maxpool4 = self.maxpool4(conv4)
+
+        center = self.center(maxpool4)
+        up4 = self.up_concat4(conv4, center)
+        up3 = self.up_concat3(conv3, up4)
+        up2 = self.up_concat2(conv2, up3)
+        up1 = self.up_concat1(conv1, up2)
+
+        final = self.final(up1)
+
+        return final