U-Net: Medical Image Segmentation

A scan that needs a map, not a verdict

During the second COVID-19 wave in India in April–May 2021, hospitals in Delhi, Mumbai, and Pune were receiving chest CT scans faster than radiologists could read them. A radiologist did not just need to know "is this patient COVID-positive" — that question was often already answered by an RT-PCR test. What the triage desk needed was a CT severity score: what fraction of the lung volume is covered by ground-glass opacities and consolidation, and where. That number decided who got an oxygen bed first. A whole-image classifier that outputs "COVID: yes, 91% confidence" is useless here — it collapses a lung's worth of information into one number. What the situation demanded was a decision for every pixel in the scan: is this pixel healthy lung tissue, is it opacity, or is it outside the lung entirely? Stack up those per-pixel decisions and you get a map a doctor can look at, measure, and trust. That task — dense, pixel-by-pixel classification — is called semantic segmentation, and U-Net, introduced by Olaf Ronneberger, Philipp Fischer, and Thomas Brox in 2015 for segmenting cells in microscopy images, is the architecture that made this practical even when a hospital or lab has only a few dozen hand-labeled training scans, not the millions of images an ImageNet-scale classifier expects.

Three questions vision models answer — and why segmentation needs different machinery

It helps to place segmentation next to the two tasks you have already studied. Classification answers "what is in this image," producing one label for the whole picture. Object detection answers "where are the objects," producing a bounding box and label per instance. Semantic segmentation answers a much stricter question: "which class does each individual pixel belong to," producing an output with (roughly) the same height and width as the input, but with a predicted class at every coordinate instead of one label for the whole frame or a handful of boxes. A bounding box around a lung opacity tells you roughly where it is; a segmentation mask tells you its exact boundary, its area in square millimetres once you know the pixel spacing, and whether it touches the pleura. That precision is exactly what a severity score, a tumour-volume measurement, or a surgical-planning tool needs. (There is a further distinction — instance segmentation separates touching objects of the same class, e.g. two adjacent cells, from each other, not just from the background. The original U-Net paper cared about exactly this, and you will see how a plain per-pixel classifier is coaxed into respecting instance boundaries later in this chapter.)

The pre-U-Net approach, and why it does not scale

Before fully convolutional architectures existed, the natural way to get a segmentation map out of a CNN — the approach used by Ciresan et al. in 2012 for neuronal membrane segmentation — was to reuse an ordinary image classifier as a patch classifier: cut out a small square patch centered on each pixel, run it through a CNN, and take the predicted class of the patch as the label for that one center pixel. Slide the window over every pixel in the image and you eventually get a full segmentation map.

This works, but the arithmetic is brutal. Suppose you use 50×50 patches with a stride of 1 pixel on a 400×400 image. The number of positions a 50-pixel-wide window can occupy along one axis of a 400-pixel image is 400 − 50 + 1 = 351, so the total number of independent forward passes needed is 351² = 123,201 — over a hundred thousand CNN evaluations for one image, and adjacent patches overlap almost completely, so nearly all of that convolutional computation is being redone from scratch for a one-pixel shift. There is also a structural trade-off baked into the patch size: a larger patch gives the network more surrounding context (useful for telling opacity from a rib shadow) but blurs exactly where the boundary is, since a big patch dilutes the influence of pixels near its edge; a smaller patch localizes precisely but starves the network of context, making it easy to confuse a small opacity with noise. You cannot fix this by picking one "correct" patch size — the two goals pull in opposite directions for a single fixed window.

Fully convolutional: one pass, every pixel, at every scale

The fix is to stop classifying patches one at a time and instead build a network that is convolutional end to end, so a single forward pass over the whole image produces labels for every pixel simultaneously — this is the fully convolutional network (FCN) idea from Long, Shelhamer, and Darrell (2015), and U-Net builds directly on it. A plain FCN, though, still has a problem: it downsamples the image through several pooling layers to build up context and a large receptive field (so the network can "see" a whole lesion, not just a few pixels of it), and then has to upsample straight back to the original resolution. But pooling is lossy — a 2×2 max-pool keeps the largest value in each block and throws away exactly where inside that block the value came from. By the time the network reaches its most abstract, most context-rich layer, it knows that there is probably an opacity somewhere in a given neighbourhood, but it has lost the precise pixel coordinates of its edge. Upsampling from there alone produces a coarse, blobby mask with mushy boundaries — technically a segmentation, but not one a radiologist could trust to measure area. U-Net's one architectural addition on top of a plain encoder-decoder FCN — the skip connection — exists specifically to repair this.

Anatomy of U-Net: two paths and a bridge

U-Net has three parts, and its name comes from the shape you get when you draw them: a contracting path down the left, a bottleneck at the bottom, and an expanding path back up the right, with horizontal connections stitching the two sides together.

The contracting path (encoder) repeats a simple block at each of several depths: two 3×3 convolutions, each followed by a ReLU, then a 2×2 max-pool with stride 2. Each pooling step halves the spatial resolution and, by convention, doubles the number of feature channels — going from a raw image to increasingly abstract feature maps that cover more of the image per "pixel" but say less about the exact geometry. This is the "what" path: it is good at recognizing that a region of tissue looks like opacity, but progressively worse at saying precisely where that region ends.

The bottleneck is the deepest, most compressed representation — the smallest spatial size, the most channels, the most abstract features, sitting at the bottom of the U.

The expanding path (decoder) mirrors the encoder in reverse. At each stage it applies a 2×2 up-convolution (a learned transposed convolution) that doubles the spatial size and halves the channel count, bringing the coarse, context-rich features back toward full resolution. This is where U-Net departs from a plain FCN decoder: before running any further convolutions, it takes the feature map from the corresponding depth of the encoder — the one with matching spatial size, captured before that information was ever pooled away — and concatenates it, channel-wise, onto the upsampled decoder features. Only after this concatenation do two more 3×3 convolutions run, fusing the coarse contextual signal (from the upsampled path) with the sharp, precisely-located signal (from the encoder skip) into one feature map. This concatenate-then-convolve step happens once at every resolution level on the way back up. A final 1×1 convolution at the top converts the last feature map's channels into per-class scores at every pixel, which a softmax turns into a probability distribution over classes for each pixel independently.

The skip connections are the whole idea. Without them you have a plain encoder-decoder that guesses roughly where a boundary is; with them, the decoder is handed the exact pre-pooling pixel evidence at the moment it needs to commit to a boundary, at every scale, not just the finest one.

Worked example: tracing every shape through the original U-Net

The original paper made a deliberate, easy-to-miss choice: every 3×3 convolution is unpadded ("valid" convolution). A valid 3×3 convolution with stride 1 shrinks an n×n input to (n−2)×(n−2), because it needs a full 3×3 neighbourhood centered on each output pixel and has no padded border to borrow from near the edges. Two such convolutions in a row shrink the map by 4 pixels total, and this shrinkage accumulates every time the signal passes through a block. Starting from the paper's 572×572 single-channel input, here is the exact shape at every stage, derived by applying "subtract 2 per 3×3 conv, halve per 2×2 max-pool, double per 2×2 up-conv" at each step:

Two things fall out of this table that a student should notice. First, the crop is not optional: at the L4 skip, the encoder feature map is 64×64 but the upsampled decoder map at that point is only 56×56, because the decoder path lost pixels on its own convolutions on the way down and up. Concatenation requires identical spatial dimensions, so the encoder map is center-cropped from 64 to 56 before the channel-wise stack — the same logic applies at all four skip levels. Second, the final output, 388×388, is smaller than the 572×572 input. That is not a bug; it is the direct, compounding cost of ten unpadded 3×3 convolutions along the longest path interleaved with four 2×2 poolings that each halve the spatial size before the next convolutions apply — see the table above for the exact, authoritative trace. To segment a full-resolution medical image with this valid-convolution design, the paper uses an overlap-tile strategy: cut the large image into overlapping 572×572 input tiles, where each tile's 388×388 output covers a distinct, non-overlapping region of the final full-size map, and use mirror-reflection padding at the true image border to synthesize the missing context outside the actual scan. Most modern re-implementations sidestep all of this by padding every convolution to keep spatial size constant ("same" padding) and simply upsampling to exactly match the encoder size, trading the historical elegance of the crop-and-tile scheme for simpler bookkeeping — that is the version used in the code below.

Code: the encoder–decoder–skip pattern, with same-padding for simplicity

The block below builds one encoder step and one matching decoder step with same-padding convolutions (kernel 3, padding 1), so spatial size only changes at the explicit pool/up-conv operations — easier to trace by hand than the original paper's valid-convolution version above.

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, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
        )
    def forward(self, x):
        return self.block(x)

enc1 = DoubleConv(1, 64)
pool = nn.MaxPool2d(kernel_size=2, stride=2)
enc2 = DoubleConv(64, 128)

x = torch.randn(1, 1, 64, 64)   # one 64x64 grayscale CT crop
f1 = enc1(x)                    # same-padding keeps spatial size: 64x64
p1 = pool(f1)                   # maxpool halves it: 32x32
f2 = enc2(p1)                   # channels double: 128, size stays 32x32

print(f1.shape, p1.shape, f2.shape)
# torch.Size([1, 64, 64, 64]) torch.Size([1, 64, 32, 32]) torch.Size([1, 128, 32, 32])

up = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
d = up(f2)                      # transposed conv doubles size back: 64x64, 64ch
merged = torch.cat([d, f1], dim=1)   # skip connection: concatenate along channel axis
dec1 = DoubleConv(128, 64)(merged)   # 64 (upsampled) + 64 (encoder f1) = 128 in

print(merged.shape, dec1.shape)
# torch.Size([1, 128, 64, 64]) torch.Size([1, 64, 64, 64])

Trace the arithmetic yourself: a 3×3 convolution with padding 1 and stride 1 leaves the spatial size unchanged, since the padded border replaces exactly what an unpadded convolution would have removed. MaxPool2d(2) halves 64 to 32. ConvTranspose2d(kernel=2, stride=2) maps a 32×32 map back to 64×64 exactly (output size = (in−1)×stride − 2×padding + kernel = 31×2 + 2 = 64). The concatenation stacks 64 upsampled channels onto 64 encoder channels to give 128, which is exactly the in_ch the next DoubleConv expects — this channel bookkeeping (upsampled channels + skip channels = next block's input channels) is the single most common source of shape-mismatch bugs when students implement U-Net from scratch.

Measuring success: why plain accuracy lies, and Dice steps in

In a chest CT slice, opacity pixels are typically a small fraction of the whole image — most of any slice is background air, bone, or unaffected lung. A model that predicts "background" for every single pixel could easily score above 95% pixel accuracy on such a scan while being clinically worthless, because it never finds the lesion at all. This is exactly the class-imbalance problem you have already met in classification, now at the pixel level, and it means accuracy is the wrong metric to optimize or report for segmentation.

The standard fix is the Dice coefficient (equivalently, the F1 score computed over the set of predicted-positive and true-positive pixels): for a predicted mask A and ground-truth mask B, Dice = 2|A∩B| / (|A| + |B|), where |A∩B| counts pixels both masks agree are positive, and |A|, |B| are simply the pixel counts of each mask. Dice is 1 when the masks match perfectly and 0 when they share no pixels, and unlike raw accuracy it is completely insensitive to how large the (usually enormous) true-negative background region is.

Work through a small case by hand. Suppose a 4×4 ground-truth mask marks a lesion at the four center cells (rows 1–2, columns 1–2, 0-indexed), so |B| = 4. A model's predicted mask correctly marks three of those four cells but misses one corner, so |A| = 3, and every one of the model's three predicted pixels is a true positive, so |A∩B| = 3. Then Dice = 2×3 / (4+3) = 6/7 ≈ 0.857. Contrast this with accuracy on the same 4×4=16-pixel grid: 15 of 16 pixels are labeled correctly (only the one missed corner is wrong), giving 93.75% accuracy — a number that looks almost perfect despite the model missing a quarter of the actual lesion. This gap is exactly why segmentation papers report Dice (or the closely related Intersection-over-Union), never raw pixel accuracy, as the headline metric.

def dice_coefficient(pred_mask, true_mask, eps=1e-7):
    intersection = (pred_mask * true_mask).sum()
    return (2 * intersection + eps) / (pred_mask.sum() + true_mask.sum() + eps)

# the 4x4 toy example above, as flat binary tensors
pred = torch.tensor([0,0,0,0, 0,1,1,0, 0,1,0,0, 0,0,0,0], dtype=torch.float32)
true = torch.tensor([0,0,0,0, 0,1,1,0, 0,1,1,0, 0,0,0,0], dtype=torch.float32)
print(dice_coefficient(pred, true).item())
# 0.8571... i.e. 6/7, matching the hand calculation

During training, the loss actually minimized is usually Dice loss (1 − Dice), often summed with a per-pixel weighted cross-entropy term. The original U-Net paper adds one more refinement worth knowing: because it needed to separate individually touching cells using only a semantic (not instance) segmentation formulation, it computes a per-pixel weight map that sharply upweights the thin background strip between two adjacent object boundaries, using a Gaussian term that grows the closer a background pixel is to two different objects at once. The effect is that a mistake on the one-pixel gap separating two touching lesions costs the network far more loss than an equally-sized mistake in open background — which is precisely what forces a per-pixel classifier to learn crisp separating boundaries it would otherwise have no explicit incentive to draw.

Common misconception: "skip connections exist to fix vanishing gradients, like ResNet"

Because U-Net's skip connections look superficially like ResNet's residual connections — both are described as "the encoder output is carried forward to a later layer" — many students conclude they solve the same problem: shortening the path gradients travel during backpropagation so early layers keep receiving a useful signal. That is not why U-Net has them, and the difference matters for understanding the architecture correctly.

ResNet's residual connections perform element-wise addition: the input to a block is added to that block's output, requiring the two to have identical shape and channel count, and the sum simply lets gradients flow through an identity shortcut around the learned layers — its purpose is training stability in very deep networks. U-Net's skip connections perform channel-wise concatenation, not addition: the encoder feature map is stacked alongside the upsampled decoder feature map along the channel axis, deliberately increasing the channel count (as you saw in the code trace above, 64+64=128), so that the following convolution has both signals available separately and can learn any combination of them it needs. The purpose is not gradient flow — it is spatial information recovery. Every max-pool in the encoder deliberately discards exact pixel position in exchange for translation invariance and a larger receptive field; by the bottleneck, the network can tell you a lesion exists somewhere in a neighbourhood but has genuinely lost the sub-pixel-precise location of its edge. The skip connection's entire job is to hand the decoder that lost high-resolution evidence back, at the exact moment it is asked to commit to where a boundary sits. Better gradient flow is a pleasant side effect of the shorter path skip connections create, but it is not the reason they are there — an encoder-decoder network can be made very deep and still fail to segment sharp boundaries without concatenated skips, because the problem being fixed is loss of location, not loss of gradient signal.

Active recall

Attempt every question before reading its answer.

1. Why can a whole-image classifier not answer "what fraction of this lung is affected by opacity," even if it is 99% accurate at detecting COVID from a CT scan?
2. A patch-based classifier uses 40×40 patches with stride 1 on a 300×300 image. How many forward passes does it need to label every pixel?
3. In the original (valid-convolution) U-Net, the input is 572×572 and the output is 388×388. Name the two operations responsible for this shrinkage, and explain why the encoder feature map must be cropped before each concatenation.
4. A ground-truth mask has 10 positive pixels. A model predicts 8 positive pixels, of which 6 correctly overlap the ground truth. Compute the Dice coefficient.
5. Why does U-Net concatenate the encoder's feature map into the decoder instead of relying only on the upsampled bottleneck features?
6. True or false: replacing U-Net's concatenation skip connections with ResNet-style element-wise addition would give an architecturally equivalent network, just with fewer channels downstream. Justify your answer.

Answers.

1. A classifier outputs one label (or one probability) for the entire image, with no notion of location or extent — it cannot distinguish "5% of the lung affected, peripheral" from "40% affected, bilateral," and both of those clinically very different scans could receive the identical "COVID-positive, 99%" label. Computing an area or a boundary requires a decision at every pixel, which is a segmentation task, not a classification task.
2. Positions along one axis: 300 − 40 + 1 = 261. Total forward passes: 261² = 68,121.
3. Every unpadded ("valid") 3×3 convolution removes a 1-pixel border on each side, shrinking an n×n map to (n−2)×(n−2); this happens twice per encoder/decoder block and accumulates over the whole network, which is why the exact-shape trace in the worked example ends at 388×388 for a 572×572 input. Cropping before concatenation is required because, by a given decoder depth, the decoder's feature map is smaller than the encoder's map at the matching depth (the decoder path has accumulated more of this valid-convolution shrinkage along its round trip), so the two cannot be stacked channel-wise until they share the same height and width.
4. Dice = 2×6 / (10 + 8) = 12/18 = 2/3 ≈ 0.667.
5. Because pooling in the encoder trades exact spatial location for context and translation invariance, so the upsampled bottleneck features alone can only produce a coarse, blurry mask — they know roughly where a region is but not its precise edge. The encoder skip carries the pre-pooling, full-resolution evidence forward so the decoder has the fine-grained information needed to draw a sharp boundary at every scale, not just a smoothed approximation of one.
6. False. Element-wise addition requires the two tensors being combined to already have the same number of channels (or a learned projection to force them to match), and it immediately collapses the fine and coarse signals into one fixed blend before any further layer sees them individually. Concatenation keeps both signals distinct and hands them, side by side, to the next convolution, which can then learn any weighting or combination of "trust the coarse context here, trust the sharp edge there" that the task needs — a flexibility a forced sum discards at the point of combination, not merely a difference in parameter count.

Think About It

Think about this: How would you explain u-net: medical image segmentation to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.

Practice Exercises

Now it is time to practice! Complete these challenges to solidify your understanding:

• Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
• Exercise 2: Find a real-world example where u-net: medical image segmentation is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
• Exercise 3: Create a mind-map connecting u-net: medical image segmentation to at least 3 other topics you have studied.

Key Takeaways — Summary and Recap

Let us recap what we covered: the core ideas behind u-net: medical image segmentation, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.