A complete image segmentation project — building a U-Net for pixel-level semantic segmentation, evaluated with IoU, extending detection's bounding boxes to precise per-pixel classification.
Problem Statement
Build a semantic segmentation model that classifies every pixel in an image into one of a small set of categories (e.g. foreground object vs background, or a few specific object types), evaluated with mean Intersection over Union (IoU).
Dataset
A dataset with pixel-level segmentation masks — the Oxford-IIIT Pet dataset (with trimap segmentation masks) is a well-suited, accessible starting point, available directly through torchvision.datasets.
Architecture & Approach
U-Net — an encoder-decoder architecture with skip connections directly linking corresponding encoder and decoder resolutions — is the standard, well-suited architecture for this task, preserving fine spatial detail that a pure encoder-decoder without skip connections would lose.
Step-by-Step Build
import torch
import torch.nn as nn
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), nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1), nn.ReLU(inplace=True)
)
def forward(self, x): return self.block(x)
class UNet(nn.Module):
def __init__(self, num_classes=3):
super().__init__()
self.enc1 = DoubleConv(3, 64)
self.enc2 = DoubleConv(64, 128)
self.pool = nn.MaxPool2d(2)
self.bottleneck = DoubleConv(128, 256)
self.upconv2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
self.dec2 = DoubleConv(256, 128) # 256 = 128 (upsampled) + 128 (skip connection)
self.upconv1 = nn.ConvTranspose2d(128, 64, 2, stride=2)
self.dec1 = DoubleConv(128, 64)
self.final = nn.Conv2d(64, num_classes, 1)
def forward(self, x):
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
b = self.bottleneck(self.pool(e2))
d2 = self.upconv2(b)
d2 = torch.cat([d2, e2], dim=1) # THE skip connection -- combines upsampled and encoder features
d2 = self.dec2(d2)
d1 = self.upconv1(d2)
d1 = torch.cat([d1, e1], dim=1)
d1 = self.dec1(d1)
return self.final(d1) # (batch, num_classes, H, W) -- one score per class, PER PIXEL
model = UNet(num_classes=3)
x = torch.randn(1, 3, 256, 256)
output = model(x)
print(output.shape) # (1, 3, 256, 256)
# Training loop -- notice the loss operates per-pixel
loss_fn = nn.CrossEntropyLoss() # target shape: (batch, H, W) with class indices per pixel
for x_batch, mask_batch in train_loader:
optimizer.zero_grad()
output = model(x_batch) # (batch, num_classes, H, W)
loss = loss_fn(output, mask_batch) # mask_batch: (batch, H, W)
loss.backward()
optimizer.step()Expected Results
With a modest U-Net and a few dozen epochs on a reasonably sized segmentation dataset, expect mean IoU in a moderate-to-good range, with segmentation boundaries visibly cleaner (and better matching the true object outline) than a decoder without skip connections would produce — worth verifying by training a version with the skip connections removed for direct comparison.
Key Learnings & Extensions
- The skip connections are the single most important architectural feature to understand here — remove them and observe how much coarser and less precise the resulting segmentation boundaries become.
- Extension: Compute Dice score in addition to IoU, and compare how the two metrics behave differently, especially on small objects.
- Extension: Visualize the predicted segmentation mask overlaid on the original image, side-by-side with the ground truth mask, for direct qualitative error analysis.
