AI & ML / Computer vision / 04_segmentation.md

Segmentation architectures

Updated 4 interview angles 5 min read source
On this page7
  1. U-Net and the encoder-decoder shape
  2. The families, and when each earns its place
  3. SAM is a different kind of tool
  4. Losses, because the default is wrong here
  5. Where it breaks in production
  6. Related
  7. Interview angle

Segmentation architectures

Segmentation predicts a class for every pixel. The architectural problem it solves is a tension: classification needs a large receptive field, which means downsampling, but per-pixel output needs full resolution. Every architecture here is an answer to how you get the resolution back.

U-Net and the encoder-decoder shape

text 
input ──┬─ down ──┬─ down ──┬─ bottleneck
        │         │         │
        skip      skip      skip   (concatenated, not added)
        │         │         │
    out ─┴── up ──┴── up ───┘

The encoder downsamples to build semantic context; the decoder upsamples back to full resolution. The skip connections are the whole idea: the decoder gets the encoder’s high-resolution feature maps concatenated in, recovering the fine boundary detail that pooling destroyed.

U-Net remains the right default for a surprising number of problems: it trains on small datasets, it is small enough to run anywhere, and it excels exactly where boundaries matter and classes are few. It came from biomedical imaging where labelled data is scarce, and those properties transfer directly to drawings and documents.

The families, and when each earns its place

Family Mechanism Good when
U-Net encoder-decoder, skips few classes, small data
DeepLab atrous convolution, ASPP multi-scale objects
SegFormer transformer, light decoder robustness, scale
Mask R-CNN detect then mask instance segmentation
SAM promptable, zero-shot annotation, interaction

Atrous (dilated) convolution deserves its own sentence because it comes up: it inflates the kernel with gaps, enlarging the receptive field without downsampling and without adding parameters. That is how DeepLab sees context while keeping resolution.

SegFormer replaces the convolutional encoder with a hierarchical transformer and pairs it with a deliberately tiny all-MLP decoder. Its selling point in practice is robustness — it degrades more gracefully on corrupted or out-of-distribution inputs than comparable CNNs, which matters when your inputs are photographs of documents rather than clean scans.

Mask R-CNN is Faster R-CNN plus a mask head on each proposal, so you get instance masks. Detectron2 is the library most people mean when they say Mask R-CNN; it is a research-oriented framework, heavier to deploy than a single-file export, and worth choosing when you need its model zoo or its extension points rather than by default.

SAM is a different kind of tool

Segment Anything is promptable: give it a point, a box or a rough mask and it returns a high-quality mask for that thing, without having been trained on your classes. It does not know what a wall is. It knows where the boundaries are.

That makes it excellent for two jobs and poor for a third:

  • Annotation acceleration. Click, get a mask, correct it. This can cut labelling time enormously, and it is the highest-value use.
  • Interactive tools. A user clicks a room, you return its polygon.
  • Not an unattended classifier. With no prompt and no class vocabulary, it produces an unlabelled soup of regions. Pairing it with a detector that supplies box prompts is the standard composition.
python 
# The common pattern: a detector supplies prompts, SAM supplies precision.
boxes = detector(image)                       # coarse, but labelled
masks = sam.predict(image, boxes=boxes)       # precise, but unlabelled
labelled = [(m, b.cls) for m, b in zip(masks, boxes)]

Losses, because the default is wrong here

Cross-entropy alone fails on imbalanced segmentation. With 3% foreground, the loss is dominated by easy background pixels and the model converges to predicting nothing.

python 
def dice_loss(pred, target, eps=1e-6):
    """pred: probabilities. Optimises overlap directly."""
    num = 2 * (pred * target).sum() + eps
    den = pred.sum() + target.sum() + eps
    return 1 - num / den

loss = 0.5 * bce(pred, target) + 0.5 * dice_loss(sigmoid(pred), target)

Dice loss optimises the overlap metric directly, which is why it handles imbalance: it is a ratio, so a small foreground does not shrink the gradient. The combination of cross-entropy plus Dice is the workhorse — cross-entropy gives stable pixel-wise gradients early, Dice shapes the boundary.

Focal loss is the other standard answer: it downweights easy examples so the hard minority dominates the gradient. Reach for it when the imbalance is extreme and Dice alone trains unstably.

Where it breaks in production

  • Boundaries are where the error is. Interior pixels are easy; a 1-2 pixel boundary band carries most of the disagreement, and it is also where annotators disagree with each other.
  • Thin structures vanish. A wall two pixels wide survives one downsample and disappears after three. Either run at higher resolution, or pick an architecture that preserves resolution (dilated convolutions rather than aggressive pooling).
  • Masks are not geometry. A mask is a raster. Anything downstream that needs lengths, areas or CAD output needs vectorisation, with its own error — see Contours, polygons and vector geometry.
  • Tiling seams. Segmenting tile by tile leaves discontinuities at the joins. Overlap and blend the logits, not the argmaxed masks.
  • Post-processing is underrated. Removing components below a minimum area, and morphological closing on the predicted mask, fixes a meaningful fraction of visible errors for almost no cost.

Interview angle 4

  • “Why does U-Net have skip connections?” - downsampling builds semantic context but destroys spatial precision. The skips hand the decoder the encoder’s high-resolution features so boundaries can be recovered. Follow-up: concatenate or add? U-Net concatenates, which keeps both signals separable at the cost of channel width.
  • “Cross-entropy is not converging on a 3%-foreground mask. What now?” - add Dice loss. It optimises overlap as a ratio, so a small foreground does not shrink the gradient. Focal loss is the alternative when imbalance is extreme.
  • “Would you use SAM for this?” - only with prompts. SAM finds boundaries, not classes, so unattended it gives you unlabelled regions. Use it to accelerate annotation, or pair it with a detector that supplies boxes.
  • “Thin walls disappear in your predictions.” - they are being pooled away. Raise input resolution, reduce downsampling depth, or use dilated convolutions to grow the receptive field without losing resolution.
Previous
Object detection