In this guide
Classification tells you what is in the image. Detection tells you where with boxes. Segmentation tells you exactly which pixels belong to what — the outline of the tumor, the precise boundary of the road, the shape of the defect. When your project needs pixel-level answers, you need segmentation, and its two landmark architectures are U-Net and Mask R-CNN.
This guide explains the three segmentation flavors, how each architecture works, when to pick which, and the training and evaluation practices specific to pixel-level tasks.
The three flavors of segmentation
| Type | Question answered | Output | Example |
|---|---|---|---|
| Semantic | What class is each pixel? | One label per pixel; instances merge | "All road pixels" in a driving scene |
| Instance | What class is each pixel, per object? | Separate mask per object instance | Each individual cell in microscopy |
| Panoptic | Both, unified | Every pixel labeled + instances separated | Complete scene parsing |
Choose by the project's question: counting objects needs instance segmentation (semantic merges touching objects into one blob); measuring area or coverage needs only semantic. Picking instance segmentation when semantic suffices doubles your labeling effort for nothing.
Note: Segmentation labels are expensive — every pixel needs a class. Before committing, confirm that bounding boxes (detection) genuinely cannot answer your project's question. Many "segmentation projects" are detection projects in disguise.
U-Net: the segmentation workhorse
U-Net (Ronneberger et al., 2015) was designed for biomedical image segmentation and became the default architecture for semantic segmentation. Its shape gives it the name: a contracting encoder path, an expanding decoder path, symmetric like a U.
- Encoder (left side): standard CNN downsampling — convolutions + pooling, doubling channels while halving resolution. Captures what is present at increasing abstraction.
- Bottleneck: the most compressed representation.
- Decoder (right side): upsampling (transposed convolutions) back to full resolution. Recovers where things are.
- Skip connections: the key innovation — feature maps from each encoder level are concatenated directly to the matching decoder level. Fine spatial detail lost in downsampling flows straight across, so the output mask aligns precisely with the input image.
Without skip connections, the decoder would have to reconstruct precise boundaries from a tiny bottleneck vector — it cannot. With them, U-Net produces sharp masks even from small training sets, which is why it dominates medical imaging where labeled data is scarce.
Modern practice swaps the encoder for a pretrained backbone (ResNet, EfficientNet) — transfer learning applies fully here (see the CNN transfer learning guide) — while keeping the U-shaped decoder with skips. Libraries like segmentation-models-pytorch give you U-Net with any backbone in a few lines:
import segmentation_models_pytorch as smp
model = smp.Unet(
encoder_name="resnet34",
encoder_weights="imagenet",
in_channels=3,
classes=2, # binary: foreground vs background
)
Mask R-CNN: detection plus masks
Mask R-CNN (He et al., 2017) extends Faster R-CNN: it detects objects with boxes and predicts a binary mask inside each box. The pipeline:
- Backbone + RPN propose regions of interest (as in Faster R-CNN).
- RoIAlign extracts fixed-size features per proposal — replacing RoIPool's quantization, which matters because mask prediction needs pixel-accurate alignment.
- Three heads run in parallel: classification, box regression, and a small FCN predicting a mask per class.
The mask head predicts a binary mask for each class independently (sigmoid per pixel, not softmax across classes) — decoupling mask prediction from classification, which the paper showed improves results.
Mask R-CNN is the choice when you need instance segmentation with the maturity of the detection ecosystem: counting, per-object measurement, scenes with overlapping instances of the same class.
U-Net vs Mask R-CNN: choosing
| Aspect | U-Net | Mask R-CNN |
|---|---|---|
| Segmentation type | Semantic | Instance |
| Labeling cost | Pixel masks (no instance separation) | Pixel masks + instance separation |
| Small datasets | Excellent (works with tens of images) | Needs more data |
| Speed | Fast (single pass, fully convolutional) | Slower (two-stage) |
| Overlapping same-class objects | Merges them | Separates them |
| Typical domains | Medical, satellite, defect inspection | People, vehicles, retail, biology counting |
| Framework support | segmentation-models-pytorch | Detectron2, torchvision |
Rule of thumb: need per-object masks (counting, tracking individuals) → Mask R-CNN. Need pixel-class maps (area, boundaries, coverage) → U-Net. When in doubt, prototype U-Net first — it trains faster and labels cheaper.
Training specifics for segmentation
Loss functions. Pixel-wise cross-entropy works, but class imbalance (99% background pixels) demands better:
| Loss | Handles | Notes |
|---|---|---|
| Cross-entropy | Baseline | Drowns in background pixels |
| Dice loss | Imbalance directly | 1 − Dice coefficient; standard for medical |
| Focal loss | Hard examples | Down-weights easy background pixels |
| Combo (Dice + CE) | Both | Common convention; start here |
Augmentation is load-bearing. With small segmentation datasets, geometric augmentation (flips, rotations, elastic deformations for medical) must transform image and mask identically — libraries like Albumentations handle the paired transforms. Photometric augmentation (brightness, contrast) applies to the image only. See data augmentation for image models.
Patch-based training for large images (whole-slide microscopy, satellite): train on cropped patches, stitch predictions at inference with overlap + blending to avoid seam artifacts.
Evaluation: IoU and Dice
Pixel accuracy is misleading (predicting all-background scores 99% on a 1% foreground task). The real metrics:
| Metric | Formula | Use |
|---|---|---|
| IoU (Jaccard) | intersection / union | Standard; penalizes both over- and under-segmentation |
| Dice (F1) | 2·intersection / (sum) | Common in medical; more forgiving of small mismatches |
| Mean IoU | Average IoU over classes | Multi-class summary |
import numpy as np
def iou(pred, target):
inter = np.logical_and(pred, target).sum()
union = np.logical_or(pred, target).sum()
return inter / union if union else 1.0
Report per-class IoU, not just the mean — a model can have a respectable mIoU while completely failing the rare class you actually care about. And evaluate on images from the deployment distribution: masks that look perfect on lab images often degrade on field images.
Common mistakes
- Choosing segmentation when detection suffices. Pixel labels cost 10–50x more than boxes. Confirm the project needs pixels.
- Pixel accuracy as the metric. The all-background trap. Use IoU/Dice.
- Augmenting image without mask. Transforms applied to only one of the pair silently corrupt training. Use paired augmentation.
- Ignoring class imbalance. Cross-entropy on 99% background learns to predict background. Use Dice/combo losses.
- Training from scratch on 50 images. Use a pretrained encoder backbone — always, for small datasets.
- Noisy boundaries in labels. Inconsistent annotation of object edges caps achievable IoU; calibrate labelers (see data labeling strategies).
- Evaluating on the training distribution only. Field images differ; test on them.
Where to go from here
Segmentation builds on the full vision pipeline: CNN transfer learning for backbones, data augmentation for small datasets, data labeling for mask annotation, and the YOLO detection guide when boxes turn out to be enough. For deploying segmentation on device, see edge AI on Raspberry Pi. The transformer architecture guide explains the attention behind newer segmentation transformers (SegFormer, Mask2Former). More vision topics in the AI & Machine Learning branch hub.