Create residual_unet.py

residual_unet.py

@@ -0,0 +1,117 @@
+import torch
+import torch.nn as nn
+from huggingface_hub import PyTorchModelHubMixin  # <-- 1. IMPORT THIS
+
+class ResidualConvBlock(nn.Module):
+    """
+    A residual convolutional block consisting of two convolutional layers,
+    batch normalization, ReLU activation, and a shortcut connection.
+    """
+    def __init__(self, in_channels, out_channels):
+        super(ResidualConvBlock, self).__init__()
+        # First convolutional layer
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        # Second convolutional layer
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        # Shortcut connection to match dimensions if in_channels != out_channels
+        self.shortcut = nn.Sequential()
+        if in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+    def forward(self, x):
+        # The output of the conv layers is added to the original input (shortcut)
+        residual = self.shortcut(x)
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += residual
+        return self.relu(out)
+
+#  ↓
+#  ↓   2. ADD THE MIXIN HERE
+#  ↓
+class ResidualUNet(nn.Module, PyTorchModelHubMixin): 
+    """
+    Residual U-Net architecture for semantic segmentation.
+    The network consists of a contracting path (encoder) and an expansive path (decoder).
+    Input size is assumed to be (448, 448, 3).
+    """
+    def __init__(self, in_channels=3, out_channels=1):
+        super(ResidualUNet, self).__init__()
+        
+        # Save arguments to config.json
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        
+        # =====================================
+        # Encoder (Contracting Path) - 5 levels
+        # =====================================
+        self.encoder1 = ResidualConvBlock(in_channels, 64)
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder2 = ResidualConvBlock(64, 128)
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder3 = ResidualConvBlock(128, 256)
+        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder4 = ResidualConvBlock(256, 512)
+        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder5 = ResidualConvBlock(512, 1024)
+        self.pool5 = nn.MaxPool2d(kernel_size=2, stride=2)
+        # =====================================
+        # Bottleneck
+        # =====================================
+        self.bottleneck = ResidualConvBlock(1024, 2048)
+        # =====================================
+        # Decoder (Expansive Path) - 5 levels
+        # =====================================
+        self.upconv5 = nn.ConvTranspose2d(2048, 1024, kernel_size=2, stride=2)
+        self.decoder5 = ResidualConvBlock(1024 + 1024, 1024) # Concatenating skip connection from encoder5
+        self.upconv4 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
+        self.decoder4 = ResidualConvBlock(512 + 512, 512) # Concatenating skip connection from encoder4
+        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
+        self.decoder3 = ResidualConvBlock(256 + 256, 256) # Concatenating skip connection from encoder3
+        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
+        self.decoder2 = ResidualConvBlock(128 + 128, 128) # Concatenating skip connection from encoder2
+        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
+        self.decoder1 = ResidualConvBlock(64 + 64, 64) # Concatenating skip connection from encoder1
+        # =====================================
+        # Output Layer
+        # =====================================
+        self.outconv = nn.Conv2d(64, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        # Encoder path
+        skip1 = self.encoder1(x)
+        p1 = self.pool1(skip1)
+        skip2 = self.encoder2(p1)
+        p2 = self.pool2(skip2)
+        skip3 = self.encoder3(p2)
+        p3 = self.pool3(skip3)
+        skip4 = self.encoder4(p3)
+        p4 = self.pool4(skip4)
+        skip5 = self.encoder5(p4)
+        p5 = self.pool5(skip5)
+        # Bottleneck
+        b = self.bottleneck(p5)
+        # Decoder path with skip connections
+        d5 = self.upconv5(b)
+        d5 = torch.cat((skip5, d5), dim=1)
+        d5 = self.decoder5(d5)
+        d4 = self.upconv4(d5)
+        d4 = torch.cat((skip4, d4), dim=1)
+        d4 = self.decoder4(d4)
+        d3 = self.upconv3(d4)
+        d3 = torch.cat((skip3, d3), dim=1)
+        d3 = self.decoder3(d3)
+        d2 = self.upconv2(d3)
+        d2 = torch.cat((skip2, d2), dim=1)
+        d2 = self.decoder2(d2)
+        d1 = self.upconv1(d2)
+        d1 = torch.cat((skip1, d1), dim=1)
+        d1 = self.decoder1(d1)
+        # Final output
+        outputs = self.outconv(d1)
+        return outputs