Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class DoubleConv(nn.Module): | |
| """ | |
| DoubleConv Module | |
| ================= | |
| A standard building block for UNet, consisting of two consecutive convolution layers. | |
| Each 3x3 convolution is followed by Batch Normalization and ReLU activation. | |
| Structure: | |
| Input -> [Conv3x3 -> BatchNorm -> ReLU] -> [Conv3x3 -> BatchNorm -> ReLU] -> Output | |
| """ | |
| def __init__(self, in_channels, out_channels): | |
| super().__init__() | |
| self.double_conv = nn.Sequential( | |
| nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(out_channels), | |
| nn.ReLU(inplace=True), | |
| nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), | |
| nn.BatchNorm2d(out_channels), | |
| nn.ReLU(inplace=True) | |
| ) | |
| def forward(self, x): | |
| return self.double_conv(x) | |
| class Down(nn.Module): | |
| """ | |
| Down Module | |
| =========== | |
| Handles the downsampling step in the encoder part of the UNet. | |
| It supports two modes of downsampling: | |
| 1. 'maxpool': Uses MaxPool2d(2) to halve the spatial dimensions. | |
| 2. 'strided': Uses a Strided Conv (kernel=3, stride=2) to halve dimensions while learning features. | |
| After downsampling, a DoubleConv block processes the features. | |
| """ | |
| def __init__(self, in_channels, out_channels, mode='maxpool'): | |
| super().__init__() | |
| self.mode = mode | |
| if mode == 'maxpool': | |
| # Option 1: MaxPool downsampling (Standard UNet) | |
| self.down_layer = nn.MaxPool2d(2) | |
| self.conv = DoubleConv(in_channels, out_channels) | |
| elif mode == 'strided': | |
| # Option 2: Strided Convolution downsampling | |
| # Replaces the pooling operation with a learnable strided convolution | |
| self.down_layer = nn.Sequential( | |
| nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1), | |
| nn.BatchNorm2d(out_channels), | |
| nn.ReLU(inplace=True) | |
| ) | |
| # The strided conv handles the channel change (in -> out). | |
| # The following DoubleConv refines these features (out -> out). | |
| self.conv = DoubleConv(out_channels, out_channels) | |
| else: | |
| raise ValueError(f"Unknown downsample mode: {mode}") | |
| def forward(self, x): | |
| if self.mode == 'maxpool': | |
| x = self.down_layer(x) | |
| return self.conv(x) | |
| else: | |
| # Strided path | |
| x = self.down_layer(x) # [B, OutCh, H/2, W/2] | |
| return self.conv(x) # [B, OutCh, H/2, W/2] | |
| class Up(nn.Module): | |
| """ | |
| Up Module | |
| ========= | |
| Handles the upsampling step in the decoder part of the UNet. | |
| It supports two modes of upsampling: | |
| 1. 'transpose': Uses ConvTranspose2d to learn how to upsample. | |
| 2. 'upsample': Uses bilinear interpolation (nn.Upsample). | |
| Steps: | |
| 1. Upsample the input tensor (x1) from the previous lower layer. | |
| 2. Concatenate it with the corresponding feature map from the encoder (x2) (Skip Connection). | |
| - Handles padding if dimensions don't match perfectly. | |
| 3. Process the combined features with a DoubleConv block. | |
| """ | |
| def __init__(self, in_channels, out_channels, mode='transpose'): | |
| super().__init__() | |
| if mode == 'transpose': | |
| # Option 1: Transpose Convolution | |
| # Typical for original UNet. Upsamples and reduces channels by half. | |
| # in_channels is the dimension of the deep feature map coming UP. | |
| self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2) | |
| self.up_mode = 'transpose' | |
| elif mode == 'upsample': | |
| # Option 2: Bilinear Upsampling | |
| # Does not reduce channels itself, so we need a 1x1 conv to reduce channels | |
| # to match the skip connection size before DoubleConv. | |
| self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) | |
| self.conv_adjust = nn.Conv2d(in_channels, in_channels // 2, kernel_size=1) | |
| self.up_mode = 'upsample' | |
| else: | |
| raise ValueError(f"Unknown upsample mode: {mode}") | |
| # DoubleConv takes the concatenated input. | |
| # Channels = (in_channels // 2 from Up) + (in_channels // 2 from Skip) = in_channels | |
| # Outputs count = out_channels | |
| self.conv = DoubleConv(in_channels, out_channels) | |
| def forward(self, x1, x2): | |
| """ | |
| x1: Input from the previous decoder layer (to be upsampled) | |
| x2: Input from the encoder layer (skip connection) | |
| """ | |
| x1 = self.up(x1) | |
| if hasattr(self, 'conv_adjust'): | |
| x1 = self.conv_adjust(x1) | |
| # Handle padding if x1 and x2 have slightly different sizes due to odd dimensions | |
| # input is CHW | |
| diffY = x2.size()[2] - x1.size()[2] | |
| diffX = x2.size()[3] - x1.size()[3] | |
| x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, | |
| diffY // 2, diffY - diffY // 2]) | |
| # Concatenate x2 (skip) and x1 (upsampled) along the channel dimension | |
| x = torch.cat([x2, x1], dim=1) | |
| return self.conv(x) | |
| class UNet(nn.Module): | |
| """ | |
| UNet Architecture | |
| ================= | |
| A U-shaped encoder-decoder architecture for image segmentation. | |
| Configurable Parameters: | |
| - n_channels: Number of input image channels (e.g., 3 for RGB). | |
| - n_classes: Number of output classes (e.g., 1 for binary mask). | |
| - downsample_mode: 'maxpool' or 'strided'. | |
| - upsample_mode: 'transpose' or 'upsample' (bilinear). | |
| """ | |
| def __init__(self, n_channels, n_classes, downsample_mode='maxpool', upsample_mode='transpose'): | |
| super(UNet, self).__init__() | |
| self.n_channels = n_channels | |
| self.n_classes = n_classes | |
| self.downsample_mode = downsample_mode | |
| self.upsample_mode = upsample_mode | |
| # Initial Feature Extraction | |
| # Input: [B, n_channels, H, W] -> Output: [B, 64, H, W] | |
| self.inc = DoubleConv(n_channels, 64) | |
| # Encoder (Downsampling Path) | |
| # Each step reduces H,W by 2 and doubles Channels | |
| # Down 1: 64 -> 128 | |
| self.down1 = Down(64, 128, mode=downsample_mode) | |
| # Down 2: 128 -> 256 | |
| self.down2 = Down(128, 256, mode=downsample_mode) | |
| # Down 3: 256 -> 512 | |
| self.down3 = Down(256, 512, mode=downsample_mode) | |
| # Bridge / Bottleneck | |
| # Standard UNet goes to 1024. | |
| self.down4 = Down(512, 1024, mode=downsample_mode) | |
| # Decoder (Upsampling Path) | |
| # Each step doubles H,W and halves Channels (logic handled in Up block) | |
| self.up1 = Up(1024, 512, mode=upsample_mode) | |
| self.up2 = Up(512, 256, mode=upsample_mode) | |
| self.up3 = Up(256, 128, mode=upsample_mode) | |
| self.up4 = Up(128, 64, mode=upsample_mode) | |
| # Final Classification Layer | |
| # Reduces 64 channels to n_classes (1 per pixel for binary) | |
| self.outc = nn.Conv2d(64, n_classes, kernel_size=1) | |
| def forward(self, x): | |
| # Encoder Path with Skip Connections | |
| x1 = self.inc(x) # [B, 64, H, W] | |
| x2 = self.down1(x1) # [B, 128, H/2, W/2] | |
| x3 = self.down2(x2) # [B, 256, H/4, W/4] | |
| x4 = self.down3(x3) # [B, 512, H/8, W/8] | |
| x5 = self.down4(x4) # [B, 1024, H/16, W/16] (Bottleneck) | |
| # Decoder Path | |
| # Pass skip connections (x4, x3, x2, x1) to Up modules | |
| x = self.up1(x5, x4) # [B, 512, H/8, W/8] | |
| x = self.up2(x, x3) # [B, 256, H/4, W/4] | |
| x = self.up3(x, x2) # [B, 128, H/2, W/2] | |
| x = self.up4(x, x1) # [B, 64, H, W] | |
| logits = self.outc(x) # [B, n_classes, H, W] | |
| return logits | |