Module 6 — Semantic segmentation: U-Net and DeepLab

Modules 2 to 5 counted objects with boxes. But the crossroads pipeline also needs to know which pixels are road, which are pavement, which are grass. Boxes cannot express that: the road is a curved surface, not a rectangle. This module builds the classifier that labels every pixel.

From an image to a pixel map

Semantic segmentation reframes classification. Instead of one class per image, the model outputs one class per pixel, producing a class map the size of the input. For our crossroads dataset with four classes (road, pavement, grass, other), a $512 \times 512$ input yields a $512 \times 512$ integer array with values in {0, 1, 2, 3}.

This is technically a per-pixel classifier, so the loss is per-pixel cross-entropy. Everything else in the module — architecture choices, loss variants, metrics — comes from one central difficulty: a normal CNN backbone downsamples the image aggressively (by 32 or more), and we need to recover the original resolution.

The encoder-decoder pattern

U-Net (2015) established the pattern. The architecture is symmetric:

• Encoder: repeated blocks of convolutions and downsampling, exactly like a classification backbone. Spatial resolution halves at each block, channel depth grows.
• Decoder: repeated blocks of upsampling and convolutions. Spatial resolution doubles, depth shrinks.
• Skip connections: the feature map from each encoder block is concatenated to the corresponding decoder block at the same resolution.

The skip connections are the key. Without them, the decoder only sees the deepest, most abstract features, whose spatial resolution has been destroyed; the resulting mask is blob-shaped. With them, the decoder can locate class boundaries at pixel accuracy while still exploiting the semantic depth of the encoder.

import torch
import torch.nn as nn

class UNetBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = 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.conv(x)

A full U-Net stacks four encoder blocks with MaxPool2d(2) between them, four decoder blocks with ConvTranspose2d(2, 2) between them, and concatenates the matching-resolution encoder output before each decoder block. A modern implementation is a handful of files; segmentation_models_pytorch provides one line:

import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",
    encoder_weights="imagenet",
    in_channels=3,
    classes=4,
)

DeepLab and atrous convolutions

U-Net solves the resolution problem by decoding. DeepLab (2015, refined through v3+) takes another route: it keeps the encoder at high resolution by replacing standard convolutions with atrous (dilated) ones.

A standard $3 \times 3$ convolution touches nine adjacent pixels. An atrous convolution with dilation rate 2 touches nine pixels spaced two apart, covering a $5 \times 5$ area with the same nine parameters. Rate 4 covers $9 \times 9$, and so on. This inflates the receptive field without downsampling and without adding parameters.

DeepLabv3+ combines the two ideas: an atrous encoder plus a small decoder. Its atrous spatial pyramid pooling module runs several rates in parallel and concatenates them, so the network sees the same location at multiple contexts at once. In practice DeepLab tends to win on urban scenes with large regions (road, sky), U-Net on medical images with fine boundaries.

Metrics for masks

Pixel accuracy is the obvious first metric and the worst possible choice on segmentation. If 90 % of the image is road, a model that labels every pixel "road" scores 90 % accuracy while being useless. Two proper metrics dominate.

Intersection over Union, per class, is the segmentation analogue of the box IoU from module 5:

$$\text{IoU}_c = \frac{|P_c \cap G_c|}{|P_c \cup G_c|}$$

where $P_c$ and $G_c$ are the sets of pixels predicted and labelled as class $c$. Mean IoU averages this across classes. Because it treats every class equally regardless of pixel count, a class present in 0.5 % of pixels weighs as much as one present in 90 %.

Dice coefficient is close to IoU numerically but rewards overlap more aggressively:

$$\text{Dice}_c = \frac{2 |P_c \cap G_c|}{|P_c| + |G_c|}$$

Dice is always at least as large as IoU. On the crossroads it is the preferred metric when small regions (pedestrians on a crossing, for instance, if we ever tried to segment them) matter, and IoU is the preferred metric for large regions where the boundary is fuzzy.

The class imbalance trap

Segmentation datasets are catastrophically imbalanced at the pixel level. A crossroads image has millions of "road" pixels and thousands of "pedestrian crossing" pixels. Standard cross-entropy averages equally over all pixels, so the model learns to output "road" everywhere and its loss drops satisfyingly to a low value.

Three fixes, applied together in practice:

1. Weight the loss by inverse class frequency. PyTorch's CrossEntropyLoss(weight=...) accepts a per-class tensor. Use the inverse of pixel counts in the training set.
2. Use a Dice or focal-Tversky loss instead of, or in addition to, cross-entropy. Dice loss is $1 - \text{Dice}$, computed differentiably.
3. Sample crops rather than full images, biased towards regions that contain rare classes. Ten crops per image, each guaranteed to touch at least one pedestrian crossing, quickly rebalances the effective distribution seen at training time.

def dice_loss(pred, target, epsilon=1e-6):
    # pred: [B, C, H, W] probabilities. target: [B, H, W] class ids.
    pred = torch.softmax(pred, dim=1)
    target_one_hot = torch.nn.functional.one_hot(target, pred.shape[1])
    target_one_hot = target_one_hot.permute(0, 3, 1, 2).float()
    intersection = (pred * target_one_hot).sum(dim=(2, 3))
    union = pred.sum(dim=(2, 3)) + target_one_hot.sum(dim=(2, 3))
    dice = (2 * intersection + epsilon) / (union + epsilon)
    return 1 - dice.mean()

A high pixel accuracy on an imbalanced set means nothing

Report mean IoU across classes, and per-class IoU separately. If the road IoU is 0.98 and the pedestrian-crossing IoU is 0.12, your model does not understand pedestrian crossings, no matter what pixel accuracy says.

Inference and post-processing

At inference the model outputs logits of shape [B, C, H, W]. argmax along the class dimension gives the label map:

model.eval()
with torch.no_grad():
    logits = model(image.unsqueeze(0).cuda())
    mask = logits.argmax(dim=1)[0].cpu().numpy()   # [H, W] of class ids

Two post-processing steps are almost always useful. Morphological opening removes speckle noise (isolated single-class pixels). Largest connected component per class keeps only the largest region, which on a crossroads eliminates spurious "road" patches on top of buildings.

Overlay the mask on the image before trusting the metric

A good mIoU with a bad-looking overlay usually points to a leaked label channel or a wrong colour map, not a modelling problem. Save an overlay for every validation image at least during the first few epochs, then again at the end. It costs nothing and catches the class of bugs that no scalar catches.

In summary

• Semantic segmentation labels every pixel; the architecture must reconcile deep-network downsampling with pixel-accurate output.
• U-Net uses an encoder-decoder with skip connections; DeepLab keeps resolution high through atrous convolutions with rate-varying spatial pyramid pooling.
• Report mean IoU across classes and per-class IoU; Dice is close to IoU and preferred for small regions; pixel accuracy alone is misleading on imbalanced data.
• Class imbalance at the pixel level is the segmentation-specific trap: fight it with class-weighted loss, Dice or focal-Tversky loss, and biased crop sampling.

Next module: instance segmentation with Mask R-CNN, which distinguishes two cars where semantic segmentation merges them.