Semantic Segmentation — Pixel-Level Classification

Colouring a map. Detection is like placing stickers on a map — one sticker per city, approximately where each city is. Segmentation is like colouring the map by region — every pixel of the country is coloured by state, every coastline is traced exactly, every river is coloured blue. Much more precise, much more useful for geography.

The challenge: to colour pixels precisely, the model needs to understand both the broad context (what is in the image) and fine spatial detail (exactly where boundaries are). Pooling layers in CNNs lose spatial detail. U-Net's skip connections restore it — that is the key architectural insight.

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

# ── U-Net from scratch — every component explicit ─────────────────────

class DoubleConv(nn.Module):
    """Two consecutive Conv2d + BN + ReLU — the basic U-Net block."""
    def __init__(self, in_ch: int, out_ch: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
        )
    def forward(self, x): return self.net(x)

class Down(nn.Module):
    """MaxPool then DoubleConv — one encoder step."""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.pool = nn.MaxPool2d(2)
        self.conv = DoubleConv(in_ch, out_ch)
    def forward(self, x): return self.conv(self.pool(x))

class Up(nn.Module):
    """Upsample, concatenate skip connection, then DoubleConv."""
    def __init__(self, in_ch, out_ch, bilinear=True):
        super().__init__()
        if bilinear:
            self.up   = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_ch, out_ch)
        else:
            self.up   = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride=2)
            self.conv = DoubleConv(in_ch, out_ch)

    def forward(self, x, skip):
        x = self.up(x)
        # Pad if sizes differ (input not perfectly divisible)
        dy = skip.size(2) - x.size(2)
        dx = skip.size(3) - x.size(3)
        x  = F.pad(x, [dx//2, dx-dx//2, dy//2, dy-dy//2])
        # Concatenate skip connection — this is the key U-Net operation
        x  = torch.cat([skip, x], dim=1)
        return self.conv(x)

class UNet(nn.Module):
    def __init__(self, in_channels: int = 3, n_classes: int = 2,
                  features: list = [64, 128, 256, 512]):
        super().__init__()
        # Encoder
        self.inc   = DoubleConv(in_channels, features[0])
        self.down1 = Down(features[0], features[1])
        self.down2 = Down(features[1], features[2])
        self.down3 = Down(features[2], features[3])
        # Bottleneck
        self.down4 = Down(features[3], features[3] * 2)
        # Decoder — in_ch = skip_ch + upsampled_ch
        self.up1   = Up(features[3] * 4, features[3])
        self.up2   = Up(features[3] * 2, features[2])
        self.up3   = Up(features[2] * 2, features[1])
        self.up4   = Up(features[1] * 2, features[0])
        # Output
        self.outc  = nn.Conv2d(features[0], n_classes, kernel_size=1)

    def forward(self, x):
        # Encoder — save 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 — receive skip connections
        x  = self.up1(x5, x4)  # concat → (B, 512, H/8, W/8)
        x  = self.up2(x,  x3)  # concat → (B, 256, H/4, W/4)
        x  = self.up3(x,  x2)  # concat → (B, 128, H/2, W/2)
        x  = self.up4(x,  x1)  # concat → (B, 64,  H,   W)

        return self.outc(x)     # (B, n_classes, H, W)

# ── Shape check ───────────────────────────────────────────────────────
model = UNet(in_channels=3, n_classes=4)  # 4 classes: bg, road, vehicle, pedestrian
x     = torch.randn(2, 3, 256, 256)
out   = model(x)

total = sum(p.numel() for p in model.parameters())
print(f"U-Net architecture:")
print(f"  Input:   {tuple(x.shape)}")
print(f"  Output:  {tuple(out.shape)}  ← same H×W as input, n_classes channels")
print(f"  Params:  {total:,}")
print(f"
Output interpretation:")
print(f"  out[b, c, h, w] = logit for class c at pixel (h, w) in batch item b")
print(f"  argmax over dim=1 → predicted class per pixel")

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.utils.data import Dataset, DataLoader

# ── Dice loss implementation ──────────────────────────────────────────
class DiceLoss(nn.Module):
    def __init__(self, smooth: float = 1.0):
        super().__init__()
        self.smooth = smooth

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """
        logits:  (B, C, H, W) — raw model output (before softmax)
        targets: (B, H, W)    — integer class indices
        """
        n_classes = logits.size(1)
        probs     = torch.softmax(logits, dim=1)

        # One-hot encode targets: (B, H, W) → (B, C, H, W)
        one_hot = torch.zeros_like(probs)
        one_hot.scatter_(1, targets.unsqueeze(1), 1)

        # Compute Dice per class, then average
        dice_per_class = []
        for c in range(n_classes):
            pred_c = probs[:, c].reshape(-1)
            true_c = one_hot[:, c].reshape(-1)
            intersection = (pred_c * true_c).sum()
            dice = (2 * intersection + self.smooth) / (
                pred_c.sum() + true_c.sum() + self.smooth
            )
            dice_per_class.append(dice)

        return 1 - torch.stack(dice_per_class).mean()

class CombinedLoss(nn.Module):
    """CrossEntropy + Dice — standard for segmentation."""
    def __init__(self, ce_weight=0.5, dice_weight=0.5,
                  class_weights=None):
        super().__init__()
        self.ce      = nn.CrossEntropyLoss(weight=class_weights)
        self.dice    = DiceLoss()
        self.ce_w    = ce_weight
        self.dice_w  = dice_weight

    def forward(self, logits, targets):
        return self.ce_w * self.ce(logits, targets) +                self.dice_w * self.dice(logits, targets)

# ── Synthetic segmentation dataset ───────────────────────────────────
class SyntheticSegDataset(Dataset):
    """Simulates a road scene segmentation dataset."""
    CLASSES = {0: 'background', 1: 'road', 2: 'vehicle', 3: 'pedestrian'}

    def __init__(self, n: int = 200, img_size: int = 128):
        self.n, self.sz = n, img_size

    def __len__(self): return self.n

    def __getitem__(self, idx):
        np.random.seed(idx)
        img  = np.random.randint(30, 220, (3, self.sz, self.sz)).astype(np.float32) / 255
        mask = np.zeros((self.sz, self.sz), dtype=np.int64)

        # Road: bottom half
        mask[self.sz//2:, :] = 1
        # Vehicles: random rectangles in road region
        for _ in range(np.random.randint(1, 4)):
            x, y = np.random.randint(0, self.sz-30), np.random.randint(self.sz//2, self.sz-20)
            mask[y:y+20, x:x+30] = 2
        # Pedestrians: small rectangles
        for _ in range(np.random.randint(0, 3)):
            x, y = np.random.randint(0, self.sz-10), np.random.randint(self.sz//3, self.sz-30)
            mask[y:y+30, x:x+10] = 3

        return torch.FloatTensor(img), torch.LongTensor(mask)

# ── Full training loop ────────────────────────────────────────────────
from torch.utils.data import random_split

dataset  = SyntheticSegDataset(n=200, img_size=128)
train_ds, val_ds = random_split(dataset, [160, 40])
train_ld = DataLoader(train_ds, batch_size=8, shuffle=True)
val_ld   = DataLoader(val_ds,   batch_size=8)

# Class weights — background is ~60% of pixels
class_weights = torch.tensor([0.5, 1.5, 3.0, 4.0])  # upweight rare classes
model     = UNet(in_channels=3, n_classes=4, features=[32, 64, 128, 256])
criterion = CombinedLoss(class_weights=class_weights)
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20)

def pixel_accuracy(logits, targets):
    preds = logits.argmax(dim=1)
    return (preds == targets).float().mean().item()

def mean_iou(logits, targets, n_classes=4):
    preds  = logits.argmax(dim=1)
    ious   = []
    for c in range(n_classes):
        pred_c = (preds == c)
        true_c = (targets == c)
        inter  = (pred_c & true_c).sum().float()
        union  = (pred_c | true_c).sum().float()
        if union > 0:
            ious.append((inter / union).item())
    return np.mean(ious) if ious else 0.0

print("Training U-Net on road scene segmentation:")
print(f"{'Epoch':>6} {'Train loss':>12} {'Val acc':>10} {'Val mIoU':>10}")
print("─" * 42)

for epoch in range(1, 21):
    model.train()
    total_loss = 0
    for imgs, masks in train_ld:
        optimizer.zero_grad()
        logits = model(imgs)
        loss   = criterion(logits, masks)
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()
    scheduler.step()

    if epoch % 4 == 0:
        model.eval()
        accs, ious = [], []
        with torch.no_grad():
            for imgs, masks in val_ld:
                logits = model(imgs)
                accs.append(pixel_accuracy(logits, masks))
                ious.append(mean_iou(logits, masks))
        print(f"  {epoch:>4}  {total_loss/len(train_ld):>12.4f}  "
              f"{np.mean(accs):>10.4f}  {np.mean(ious):>10.4f}")

import torch
import numpy as np

def segmentation_metrics(pred_mask: torch.Tensor,
                          true_mask: torch.Tensor,
                          n_classes: int,
                          ignore_index: int = 255) -> dict:
    """
    Compute comprehensive segmentation metrics.
    pred_mask: (H, W) — predicted class indices
    true_mask: (H, W) — ground truth class indices
    """
    # Ignore unlabelled pixels (index 255 used in many datasets)
    valid = true_mask != ignore_index
    pred  = pred_mask[valid]
    true  = true_mask[valid]

    # Pixel accuracy
    pixel_acc = (pred == true).float().mean().item()

    # Per-class IoU
    class_ious = {}
    for c in range(n_classes):
        pred_c = (pred == c)
        true_c = (true == c)
        inter  = (pred_c & true_c).sum().float()
        union  = (pred_c | true_c).sum().float()
        if union > 0:
            class_ious[c] = (inter / union).item()

    # Frequency-weighted IoU — weights by class pixel frequency
    freq_iou = 0.0
    total_pixels = valid.sum().float()
    for c, iou in class_ious.items():
        freq = (true == c).sum().float() / total_pixels
        freq_iou += freq * iou

    return {
        'pixel_accuracy': pixel_acc,
        'mean_iou':       np.mean(list(class_ious.values())),
        'freq_iou':       freq_iou.item(),
        'class_iou':      class_ious,
    }

# ── Demonstrate metric sensitivity to class imbalance ─────────────────
H, W = 256, 256
n_classes = 4

# Scenario: 80% background, 15% road, 4% vehicle, 1% pedestrian
true_mask = torch.zeros(H, W, dtype=torch.long)
true_mask[H//4:,    :]           = 1  # road (bottom 75%)
true_mask[H//2:H*3//4, W//4:W//2] = 2  # vehicles
true_mask[H//3:H//3+20, W//3:W//3+10] = 3  # pedestrians

# Model that predicts everything as background
all_background = torch.zeros_like(true_mask)

# Model that predicts perfectly
perfect = true_mask.clone()

for name, pred in [('All background', all_background), ('Perfect', perfect)]:
    m = segmentation_metrics(pred, true_mask, n_classes)
    print(f"{name}:")
    print(f"  Pixel accuracy: {m['pixel_accuracy']:.4f}  ← misleading for all-bg!")
    print(f"  Mean IoU:       {m['mean_iou']:.4f}  ← correctly shows all-bg is bad")
    for c, iou in m['class_iou'].items():
        classes = ['background', 'road', 'vehicle', 'pedestrian']
        print(f"    {classes[c]:<12}: IoU = {iou:.4f}")
    print()

print("Lesson: Always report mIoU, not just pixel accuracy.")

# pip install transformers torch torchvision

from transformers import (
    SegformerForSemanticSegmentation,
    SegformerImageProcessor,
)
import torch
import numpy as np
from PIL import Image

# ── Load pretrained SegFormer ─────────────────────────────────────────
# b0=lightest, b5=heaviest. b2 is a good production balance.
model_name = 'nvidia/segformer-b0-finetuned-ade-512-512'
processor  = SegformerImageProcessor.from_pretrained(model_name)
model      = SegformerForSemanticSegmentation.from_pretrained(model_name)
model.eval()

print(f"SegFormer-b0:")
print(f"  Parameters: {sum(p.numel() for p in model.parameters()):,}")
print(f"  Classes:    {model.config.num_labels} (ADE20K dataset)")

# ── Inference ─────────────────────────────────────────────────────────
img = Image.fromarray(
    np.random.randint(50, 200, (512, 512, 3), dtype=np.uint8)
)
inputs  = processor(images=img, return_tensors='pt')
with torch.no_grad():
    outputs = model(**inputs)

# SegFormer outputs at 1/4 resolution — upsample to original size
logits = outputs.logits          # (1, 150, 128, 128)
upsampled = torch.nn.functional.interpolate(
    logits, size=img.size[::-1],  # (H, W)
    mode='bilinear', align_corners=False,
)
pred_mask = upsampled.argmax(dim=1).squeeze().numpy()
print(f"
Prediction mask shape: {pred_mask.shape}")
print(f"Unique classes predicted: {np.unique(pred_mask).tolist()[:10]}")

# ── Fine-tuning SegFormer on custom classes ───────────────────────────
print("""
# Fine-tuning SegFormer for custom segmentation:

from transformers import SegformerForSemanticSegmentation, TrainingArguments, Trainer

# 1. Load with custom number of classes
model = SegformerForSemanticSegmentation.from_pretrained(
    'nvidia/segformer-b2-finetuned-ade-512-512',
    num_labels=4,                    # your classes
    id2label={0:'bg', 1:'road', 2:'vehicle', 3:'pedestrian'},
    label2id={'bg':0, 'road':1, 'vehicle':2, 'pedestrian':3},
    ignore_mismatched_sizes=True,    # replaces the head for your classes
)

# 2. Dataset returns: pixel_values (3,H,W) + labels (H,W) integer mask
# 3. TrainingArguments — same as classification but with eval_do_concat=False
# 4. Custom compute_metrics using mIoU
# 5. trainer.train()

# Key difference from classification fine-tuning:
# - Labels are 2D masks not 1D class vectors
# - Loss is CrossEntropy over all pixels (model handles this internally)
# - Evaluation uses mIoU not accuracy
""")

# ── torchvision pretrained segmentation models ─────────────────────────
import torchvision.models.segmentation as seg_models

# DeepLabV3 with ResNet50 backbone — pretrained on COCO
deeplab = seg_models.deeplabv3_resnet50(pretrained=False, num_classes=4)
# pretrained=True in production — downloads COCO weights

x   = torch.randn(2, 3, 256, 256)
out = deeplab(x)
print(f"
DeepLabV3 output:")
print(f"  'out' key shape:  {tuple(out['out'].shape)}  ← main prediction")
print(f"  'aux' key shape:  {tuple(out['aux'].shape)}  ← auxiliary loss head")

import numpy as np
from scipy import ndimage

def clean_segmentation_mask(raw_mask: np.ndarray,
                             min_region_size: int = 50) -> np.ndarray:
    """
    Post-process a raw per-pixel argmax prediction before it ships.
    A model's raw output is almost never this clean — it typically has
    small, spurious blobs of the wrong class scattered across the image,
    one or two pixels of noise at every real boundary, and small holes
    inside otherwise-solid regions.
    """
    cleaned = raw_mask.copy()

    for class_id in np.unique(raw_mask):
        if class_id == 0:   # skip background
            continue

        class_mask = (raw_mask == class_id)

        # Label each separate connected blob of this class
        labelled, n_blobs = ndimage.label(class_mask)

        # Drop any blob smaller than min_region_size — this removes the
        # scattered single-pixel and few-pixel noise a raw model produces
        # without touching the real, large regions
        sizes = ndimage.sum(class_mask, labelled, range(1, n_blobs + 1))
        for blob_id, size in enumerate(sizes, start=1):
            if size < min_region_size:
                cleaned[labelled == blob_id] = 0

        # Fill small holes inside otherwise-solid regions of this class
        filled = ndimage.binary_fill_holes(cleaned == class_id)
        cleaned[filled] = class_id

    return cleaned

# ── Simulate a noisy raw prediction ───────────────────────────────────
np.random.seed(42)
H, W = 128, 128
raw = np.zeros((H, W), dtype=np.int64)
raw[H//2:, :] = 1                       # road region

# Scatter small spurious "vehicle" blobs — the kind of noise a real
# model produces at low-confidence pixels near boundaries
noise_points = np.random.randint(0, H, (40, 2))
for y, x in noise_points:
    raw[max(0,y-1):y+1, max(0,x-1):x+1] = 2

n_blobs_before = ndimage.label(raw == 2)[1]
cleaned = clean_segmentation_mask(raw, min_region_size=50)
n_blobs_after = ndimage.label(cleaned == 2)[1]

print(f"Spurious 'vehicle' blobs before cleanup: {n_blobs_before}")
print(f"Spurious 'vehicle' blobs after cleanup:  {n_blobs_after}")
print("The mask that actually ships to a downstream system is the")
print("cleaned one — connected-component filtering and hole-filling are")
print("standard, near-universal steps between model output and production use.")

Detection localises individual object instances — even improved to pixel precision, its output is still "here is object #1, here is object #2." Semantic segmentation has no concept of instances at all: it labels every pixel by class, so two cars standing next to each other both get painted the same "vehicle" colour and merge into one region in the mask. If you need to count how many cars are present or track a specific one, you need instance segmentation (or panoptic segmentation, which combines both) — semantic segmentation genuinely cannot answer "which one," only "what."

Downsampling in the encoder is not just a cost-saving shortcut — pooling is what lets each successive layer's receptive field cover a larger area of the original image, which is how the network builds up context about what it's looking at rather than just raw local pixel patterns. A network that never downsampled would struggle to tell "this patch of pixels is part of a road" from "this patch of pixels is part of a wall," because it would never see enough surrounding context. The actual problem downsampling causes is losing precise spatial detail (where exactly the boundary is) — which is what U-Net's skip connections are specifically designed to recover, rather than avoiding downsampling altogether.

They look similar — both copy information from an earlier layer forward — but they solve different problems. ResNet's skip connections give gradients a shortcut path so very deep classification networks stay trainable. U-Net's skip connections concatenate the encoder's feature map directly into the decoder at the matching resolution, so the decoder has access to the fine spatial detail (exact edges, precise boundaries) that pooling destroyed on the way down. The bottleneck alone tells the decoder *what* is in the image; the skip connections tell it *exactly where* the boundaries are. Removing them doesn't primarily hurt training stability — it produces blurry, imprecise mask boundaries.

In classification, class imbalance across the dataset is a data problem you can address with resampling. In segmentation, the imbalance is baked into every single image — most pixels in a road scene are background or road, and only a small fraction belong to the object you actually care about (a pedestrian, a defect, a tumour). Plain cross-entropy treats every pixel equally, so a model can achieve very low loss by getting the abundant background pixels right and largely ignoring the rare foreground class. That's exactly why segmentation commonly uses class-weighted cross-entropy, Dice loss, or a combination of both — Dice directly measures overlap and isn't dominated by whichever class has the most pixels.

Pixel accuracy just counts what fraction of pixels got the right label, and with severe class imbalance that number can be high for a model that has essentially learned nothing useful. If 90% of an image's pixels are background, a model that predicts "background" everywhere scores 90% pixel accuracy while completely failing at the task — it never gets a single foreground pixel right. Mean IoU (mIoU), which averages IoU per class rather than counting pixels in aggregate, exposes this immediately: the all-background model scores near-zero IoU on every foreground class, which is why mIoU (not pixel accuracy) is the metric that actually gets reported in segmentation papers and production dashboards.

Semantic segmentation labels every pixel by class only — all cars share one colour, with no distinction between individual cars. Instance segmentation goes further and gives each object instance its own mask, so car #1 and car #2 are separate, but it typically leaves background ("stuff" like sky and road) unlabelled. Panoptic segmentation combines both: every pixel gets a class, and every "thing" (countable object) also gets an instance ID. I'd use semantic segmentation for something like measuring total road surface area, instance segmentation for counting or tracking individual objects (how many vehicles, which one is which across frames), and panoptic segmentation when a system needs complete scene understanding, like autonomous driving.

The encoder progressively downsamples the image through pooling, which grows the receptive field so the network can understand broad context — what kind of scene this is, what objects are present. But pooling throws away precise spatial information along the way. The decoder then upsamples back to the original resolution to produce a full per-pixel mask. Skip connections concatenate each encoder stage's feature map directly into the corresponding decoder stage at the same spatial resolution, handing the decoder back the fine detail — exact edges, precise boundaries — that pooling removed. Without skip connections, U-Net would still classify regions roughly correctly but produce blurry, imprecise boundaries; the skip connections are what make its masks sharp.

Pixel accuracy is dominated by whichever class has the most pixels, which in most real-world segmentation tasks is the background. A model that predicts background for every pixel can score 90%+ pixel accuracy on a scene that's 90% background while completely failing to segment the object that actually matters. IoU (intersection over union between predicted and true regions) and Dice score both measure per-class overlap directly, so a class with very few pixels still gets fairly evaluated — the all-background model would score near zero on that class's IoU, immediately exposing the failure that pixel accuracy hides. Reporting mean IoU across classes is the standard because it weights every class's correctness equally rather than by pixel count.

A few complementary approaches, usually combined: weight the cross-entropy loss inversely by class pixel frequency so rare classes contribute proportionally more to the loss; use Dice loss (or a combined CrossEntropy + Dice loss), since Dice measures overlap directly and isn't dominated by whichever class has the most pixels; and, at the data level, oversample images that contain more of the rare class, or crop training patches centred on foreground regions rather than random crops that are mostly background. I'd also verify the fix worked by checking per-class IoU during training, not just the aggregate loss — the loss curve looking fine can hide a model that has quietly given up on the rare class.

Training U-Net from scratch needs a reasonably sized labelled dataset because every weight, including the low-level feature detectors, has to be learned from your data alone — that's thousands of labelled masks in most real settings. Fine-tuning a pretrained model like SegFormer or DeepLabV3+ reuses an ImageNet-pretrained backbone (and often a segmentation head pretrained on ADE20K or Cityscapes), so I only need to adapt the final classifier to my class set, which works with far fewer labelled images. In practice: swap in the new number of output classes, use a much smaller learning rate on the backbone than the head, and remember architecture-specific quirks — SegFormer, for instance, outputs logits at 1/4 resolution, so I need to upsample before computing the final per-pixel prediction or the loss.