Create unet.py

unet.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+
+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)