U-Net and Variants for Medical Image Analysis

Module: Computer Vision | Difficulty: Premium

U-Net Architecture

Encoder-decoder with skip connections for precise localization:

Dice Loss

Focal Tversky Loss

Boundary Loss

where is the signed distance transform.

import torch
import torch.nn as nn
import torch.nn.functional as F

class AttentionUNet(nn.Module):
    def __init__(self, in_ch=1, num_classes=1):
        super().__init__()
        self.enc1 = DoubleConv(in_ch, 64)
        self.enc2 = DoubleConv(64, 128)
        self.enc3 = DoubleConv(128, 256)
        self.pool = nn.MaxPool2d(2)
        self.attn3 = AttentionGate(256, 128, 64)
        self.attn2 = AttentionGate(128, 64, 64)
        self.up3 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.dec3 = DoubleConv(256, 128)
        self.dec2 = DoubleConv(128, 64)
        self.out = nn.Conv2d(64, num_classes, 1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        a3 = self.attn3(e3, e2)
        d3 = self.dec3(torch.cat([self.up3(e3), a3], dim=1))
        a2 = self.attn2(d3, e1)
        d2 = self.dec2(torch.cat([self.up2(d3), a2], dim=1))
        return self.out(d2)

class AttentionGate(nn.Module):
    def __init__(self, g_ch, x_ch, inter_ch):
        super().__init__()
        self.W_g = nn.Conv2d(g_ch, inter_ch, 1, bias=False)
        self.W_x = nn.Conv2d(x_ch, inter_ch, 1, bias=False)
        self.psi = nn.Sequential(
            nn.Conv2d(inter_ch, 1, 1, bias=False),
            nn.Sigmoid())
        self.relu = nn.ReLU(inplace=True)

    def forward(self, g, x):
        g1 = self.W_g(g)
        x1 = self.W_x(x)
        psi = self.relu(g1 + x1)
        psi = self.psi(psi)
        return x * psi

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
        )
    def forward(self, x):
        return self.block(x)

Research Insight: nnU-Net demonstrated that a single, well-configured U-Net architecture can match or exceed highly specialized models across 23 medical segmentation tasks. The key insight is that architecture design matters less than training procedure: adaptive preprocessing, self-configuring hyperparameters, and ensemble inference. Recent work on promptable segmentation (SAM-Med3D) enables zero-shot medical segmentation from text or point prompts, potentially eliminating the need for task-specific training in clinical deployment.