{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": "# Lesson 41: Semantic Segmentation\n\nClassification labels a whole image. Detection (Lessons 39-40) labels a handful of boxes. **Semantic segmentation** goes one step further: label *every pixel* with a class. This lesson builds a small **U-Net**-style encoder-decoder (Ronneberger, Fischer & Brox, 2015★) from scratch, and shows concretely why the architecture's defining feature — skip connections between matching encoder and decoder resolutions — matters, not just as a Lesson-35-style gradient-flow trick but for preserving spatial detail that pooling destroys."
},
{
"cell_type": "code",
"id": "b259f7a5",
"source": "import numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport matplotlib.pyplot as plt",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "fbcfd6d4",
"source": "## A 3-class pixel labeling task\n\nEach scene has several small circles and squares scattered on a background. The target isn't one label per image — it's a full per-pixel class map: 0 = background, 1 = circle, 2 = square.",
"metadata": {}
},
{
"cell_type": "code",
"id": "4042a733",
"source": "def make_scene(rng, size=32, n_shapes=5):\n scene = np.zeros((size, size), dtype=np.float32)\n mask = np.zeros((size, size), dtype=np.int64) # 0=background, 1=circle, 2=square\n for _ in range(n_shapes):\n shape_type = rng.choice([1, 2])\n r = rng.integers(2, 4)\n cx, cy = rng.integers(r, size - r), rng.integers(r, size - r)\n yy, xx = np.mgrid[0:size, 0:size]\n if shape_type == 1:\n m = ((xx - cx) ** 2 + (yy - cy) ** 2) <= r ** 2\n else:\n m = (np.abs(xx - cx) <= r) & (np.abs(yy - cy) <= r)\n scene[m] = 1.0\n mask[m] = shape_type\n scene = np.clip(scene + rng.normal(0, 0.05, scene.shape), 0, 1).astype(np.float32)\n return scene, mask\n\nrng = np.random.default_rng(13)\nN = 300\nscenes, masks = [], []\nfor _ in range(N):\n s, m = make_scene(rng)\n scenes.append(s); masks.append(m)\nscenes = np.array(scenes, dtype=np.float32)\nmasks = np.array(masks, dtype=np.int64)\n\nsplit = int(0.85 * N)\nXtr, Mtr = scenes[:split], masks[:split]\nXte, Mte = scenes[split:], masks[split:]\n\nprint('class pixel fractions (train):', dict(zip(['background', 'circle', 'square'],\n (np.bincount(Mtr.ravel()) / Mtr.size).round(3))))\n\nfig, axes = plt.subplots(2, 4, figsize=(9, 4.5))\nfor i in range(4):\n axes[0, i].imshow(Xtr[i], cmap='gray'); axes[0, i].axis('off')\n axes[1, i].imshow(Mtr[i], cmap='viridis', vmin=0, vmax=2); axes[1, i].axis('off')\naxes[0, 0].set_title('input', fontsize=9, loc='left')\naxes[1, 0].set_title('per-pixel label', fontsize=9, loc='left')\nplt.show()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "5bda9764",
"source": "## An encoder-decoder with skip connections\n\nA segmentation network needs an output the same spatial size as the input, but with a class-probability vector at every pixel instead of one RGB value. The **encoder** half is an ordinary CNN, downsampling twice via max pooling (Lesson 33) to build up wide-receptive-field, semantically rich features (Lesson 11's pyramid, again). The **decoder** half upsamples back to full resolution. A plain encoder-decoder would only have the coarse, heavily-pooled bottleneck features to work with when reconstructing full resolution — so at each decoder stage, this network also concatenates in the encoder's feature map from the *matching* resolution, before pooling destroyed that detail. This is the \"skip connection\" that gives U-Net its name (and its U-shaped diagram) — architecturally similar to Lesson 35's residual connections, but concatenating features across the encoder/decoder divide rather than adding a delta within a single stack.",
"metadata": {}
},
{
"cell_type": "code",
"id": "872a090e",
"source": "class UNetTiny(nn.Module):\n def __init__(self, n_classes=3):\n super().__init__()\n self.enc1 = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU())\n self.enc2 = nn.Sequential(nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())\n self.enc3 = nn.Sequential(nn.Conv2d(32, 64, 3, padding=1), nn.ReLU())\n self.pool = nn.MaxPool2d(2)\n self.up = nn.Upsample(scale_factor=2, mode='nearest')\n self.dec2 = nn.Sequential(nn.Conv2d(64 + 32, 32, 3, padding=1), nn.ReLU())\n self.dec1 = nn.Sequential(nn.Conv2d(32 + 16, 16, 3, padding=1), nn.ReLU())\n self.out = nn.Conv2d(16, n_classes, 1)\n\n def forward(self, x):\n f1 = self.enc1(x) # (B,16,H,W)\n f2 = self.enc2(self.pool(f1)) # (B,32,H/2,W/2)\n f3 = self.enc3(self.pool(f2)) # (B,64,H/4,W/4)\n d2 = self.dec2(torch.cat([self.up(f3), f2], dim=1)) # skip from f2\n d1 = self.dec1(torch.cat([self.up(d2), f1], dim=1)) # skip from f1\n return self.out(d1)\n\ntorch.manual_seed(0)\nmodel = UNetTiny()\nopt = torch.optim.Adam(model.parameters(), lr=0.01)\nXt = torch.tensor(Xtr).unsqueeze(1); Mt = torch.tensor(Mtr)\nfor _ in range(200):\n opt.zero_grad()\n loss = F.cross_entropy(model(Xt), Mt)\n loss.backward()\n opt.step()\n\nwith torch.no_grad():\n preds = model(torch.tensor(Xte).unsqueeze(1)).argmax(1).numpy()\n\ndef mean_iou(preds, targets, n_classes=3):\n ious = []\n for c in range(n_classes):\n p, t = preds == c, targets == c\n inter, union = (p & t).sum(), (p | t).sum()\n ious.append(inter / union if union > 0 else float('nan'))\n return ious\n\npixel_acc = (preds == Mte).mean()\nious = mean_iou(preds, Mte)\nprint(f'pixel accuracy: {pixel_acc:.1%}')\nfor name, iou in zip(['background', 'circle', 'square'], ious):\n print(f' {name:>10} IoU: {iou:.3f}')\nprint(f'mean IoU: {np.mean(ious):.3f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "295c15c0",
"source": "## Do the skip connections actually matter?\n\nTrain an otherwise-identical network with the skip connections removed — the decoder only ever sees the pooled, upsampled bottleneck features, never the original-resolution encoder features.",
"metadata": {}
},
{
"cell_type": "code",
"id": "8078efd5",
"source": "class UNetNoSkip(nn.Module):\n def __init__(self, n_classes=3):\n super().__init__()\n self.enc1 = nn.Sequential(nn.Conv2d(1, 16, 3, padding=1), nn.ReLU())\n self.enc2 = nn.Sequential(nn.Conv2d(16, 32, 3, padding=1), nn.ReLU())\n self.enc3 = nn.Sequential(nn.Conv2d(32, 64, 3, padding=1), nn.ReLU())\n self.pool = nn.MaxPool2d(2)\n self.up = nn.Upsample(scale_factor=2, mode='nearest')\n self.dec2 = nn.Sequential(nn.Conv2d(64, 32, 3, padding=1), nn.ReLU())\n self.dec1 = nn.Sequential(nn.Conv2d(32, 16, 3, padding=1), nn.ReLU())\n self.out = nn.Conv2d(16, n_classes, 1)\n\n def forward(self, x):\n f1 = self.enc1(x)\n f2 = self.enc2(self.pool(f1))\n f3 = self.enc3(self.pool(f2))\n d2 = self.dec2(self.up(f3)) # no skip: only the pooled, upsampled features\n d1 = self.dec1(self.up(d2)) # no skip\n return self.out(d1)\n\ntorch.manual_seed(0)\nmodel_noskip = UNetNoSkip()\nopt2 = torch.optim.Adam(model_noskip.parameters(), lr=0.01)\nfor _ in range(200):\n opt2.zero_grad()\n loss = F.cross_entropy(model_noskip(Xt), Mt)\n loss.backward()\n opt2.step()\n\nwith torch.no_grad():\n preds_noskip = model_noskip(torch.tensor(Xte).unsqueeze(1)).argmax(1).numpy()\n\npixel_acc_noskip = (preds_noskip == Mte).mean()\nious_noskip = mean_iou(preds_noskip, Mte)\n\nprint(f'{\"\":>18} {\"pixel acc\":>10} {\"mean IoU\":>10}')\nprint(f'{\"no skip\":>18} {pixel_acc_noskip:>10.1%} {np.mean(ious_noskip):>10.3f}')\nprint(f'{\"with skip\":>18} {pixel_acc:>10.1%} {np.mean(ious):>10.3f}')",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"id": "089cdb00",
"source": "fig, axes = plt.subplots(4, 4, figsize=(9, 9))\nfor i in range(4):\n axes[0, i].imshow(Xte[i], cmap='gray')\n axes[1, i].imshow(Mte[i], cmap='viridis', vmin=0, vmax=2)\n axes[2, i].imshow(preds[i], cmap='viridis', vmin=0, vmax=2)\n axes[3, i].imshow(preds_noskip[i], cmap='viridis', vmin=0, vmax=2)\n for r in range(4):\n axes[r, i].axis('off')\nfor r, name in enumerate(['input', 'true mask', 'pred (skip)', 'pred (no skip)']):\n axes[r, 0].set_title(name, fontsize=9, loc='left')\nplt.tight_layout()\nplt.show()",
"metadata": {},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"id": "8ec35daf",
"source": "With these small, closely-packed shapes, removing the skip connections costs several points of mean IoU. The reason is exactly what the architecture predicts: after two rounds of pooling, the bottleneck has only a quarter of the spatial resolution, and small shapes (radius 2-3 pixels) can blur together or lose their boundaries entirely at that resolution. Upsampling a blurry, low-resolution guess doesn't recover the lost detail — nearest-neighbor upsampling (or any fixed interpolation) can only spread existing information around, not invent missing edges. The skip connection sidesteps the problem by handing the decoder the original-resolution features directly, so precise boundaries never had to survive the bottleneck in the first place.\n\n### Exercise\n\n1. Increase `n_shapes` from 5 to 10, making the scene more crowded. Does the skip-vs-no-skip gap in mean IoU get larger or smaller? What does that suggest about when skip connections matter most?\n2. This lesson's loss is plain per-pixel cross-entropy. Print the per-pixel class weights implied by `class pixel fractions` above and try `F.cross_entropy(logits, Mt, weight=inverse_class_freq)` (Lesson 37's imbalance fix, applied here) to see whether it changes the circle/square IoU balance.\n3. `nn.Upsample(mode='nearest')` was used for simplicity. Try `mode='bilinear', align_corners=False` instead (Lesson 9's bilinear interpolation, now inside a network) and compare mean IoU for both the skip and no-skip models. Does the smoother upsampling help more or less than the skip connection does?",
"metadata": {}
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}