SegNet Semantic Segmentation
Pixel Labelling with Pooling-Index Unpooling
Give it a photo and it labels every pixel: road, car, traffic sign, pet or background.
- 29.5 M parameters, as in the paper
- 3 datasets
Built with
PyTorchtorchvision
pytestMatplotlib
NumPy
Kaggle
Google Colab
Overview
Semantic segmentation gives every pixel a class. SegNet does it with an encoder-decoder: the 13 convolutional layers of VGG16 shrink the image 32 times, then a mirror decoder brings it back to full size. Its particular idea is how the decoder upsamples. Each max-pool in the encoder also records which of the four pixels in every window held the maximum, and the decoder puts each value back at that position. Boundaries stay sharp and the decoder stores a few bits per window instead of whole feature maps.
I implemented it in plain PyTorch with no pretrained weights and no segmentation library: the VGG16-BN encoder, the index-based max-unpooling decoder, a cross-entropy plus soft dice loss that ignores unlabelled pixels, pixel accuracy and mIoU from a confusion matrix over the whole validation set. The width is one number, so the same code is the paper's 29.5M-parameter model on a GPU or a 1.8M-parameter one on a CPU.
The project started in 2023 as a TensorFlow notebook on Mapillary Vistas traffic signs. Porting it exposed several bugs: the training set had been overwritten by the validation set, palette label PNGs were decoded as grayscale, the image and its mask were flipped independently, and the loss was MSE on class ids. The package fixes each one and keeps the notebook in a legacy folder with the list. The test suite verifies that a small SegNet memorises four shape images; numbers from longer runs are not yet published.
The question
What does SegNet's pooling-index unpooling actually do, and can a from-scratch implementation with no pretrained weights learn to label every pixel on toy and real images?
Key features
- VGG16 encoder: 13 conv → batch-norm → ReLU layers in 5 stages, each ending in a 2 × 2 max-pool that also returns its indices.
- Mirror decoder: 5 stages of max-unpooling with the matching encoder indices followed by convolutions, then a 1 × 1 conv to one logit per class, so any input size from 32 × 32 up works.
- One width knob: base_channels 64 gives the paper's 29.5M parameters, 32 gives 7.4M for a small GPU, 16 gives 1.8M for a CPU; optional dropout on the deepest map as in Bayesian SegNet.
- Loss and metrics in one file: cross-entropy plus a weighted soft dice term for small classes, pixel accuracy and mIoU, with label 255 ignored everywhere.
- Three datasets through one flag: on-the-fly shapes, Oxford-IIIT Pet downloaded automatically, and Mapillary Vistas with all 65 classes or a chosen subset such as traffic-sign front and back.
- Training loop with AdamW, warmup plus cosine schedule, gradient clipping, float16 autocast on CUDA, best-mIoU checkpointing and resume for interrupted Kaggle sessions.
- A pytest suite: output size equals input size, the paper's parameter count, at most one non-zero per unpooled window, palette label decoding, and an overfit sanity check on four images.
How it works
Encode and remember the maxima
Five stages of conv, batch-norm and ReLU with a 2 × 2 max-pool after each one reduce an H × W image to H/32 × W/32; every pool also returns which of its four pixels held the max.
Unpool with the indices
Each decoder stage places its values back at the recorded positions, leaving the other three positions of each window at zero, and convolutions fill in the sparse map.
One logit per class
A final 1 × 1 convolution gives num_classes logits per pixel; training uses softmax cross-entropy plus an optional soft dice term, and prediction takes the argmax.
Measure on the whole split
A confusion matrix is accumulated over every validation pixel; pixel accuracy and mean IoU come from it, and the checkpoint with the best mIoU is kept.
Pick the data
Shapes check the pipeline in minutes, Oxford-IIIT Pet is the first real photos, and Mapillary Vistas palette labels are mapped through a lookup table to all classes or a chosen subset.
Examples
Input
segnet-train --dataset shapes --image-size 64 --base-channels 16 --shapes-train-size 1000 --epochs 12 --lr 3e-3 --output-dir checkpoints/shapes
Output
checkpoints/shapes/: best.pt (best validation mIoU, with configs and class names), last.pt, history.json (per-epoch train / val loss, pixel accuracy, mIoU, lr), history.png and predictions.png (image / ground truth / prediction for 4 validation images).
Input
segnet-predict --checkpoint checkpoints/pets/best.pt --images cat.jpg dog.jpg
Output
outputs/cat_mask.png, outputs/cat_overlay.png, outputs/dog_mask.png, outputs/dog_overlay.png
Tech stack
Encoder, index-based unpooling decoder, loss, metrics and training loop
Oxford-IIIT Pet download and image transforms
Model, data and metric checks plus the overfit sanity test
Training curves and prediction grids
Shape generation and label lookup tables
GPU training notebook for Pet and Mapillary Vistas
The same notebook
What I learned
- Unpooling with the saved indices puts every value back where its max came from, so boundaries stay sharp without storing whole encoder feature maps like U-Net does
- The 2023 notebook trained on its validation split, decoded palette labels as grayscale and regressed class ids with MSE; the port only meant something once those were fixed and real metrics were added
Limitations
- No mIoU or pixel-accuracy numbers from full training runs are published in the repo; the only verified figure is the overfit test
- Trains from random initialisation: the VGG16 ImageNet pretraining of the paper is not used
Keywords
SegNet · semantic segmentation · encoder-decoder · max-unpooling · pooling indices · VGG16 · PyTorch · mIoU · dice loss · Oxford-IIIT Pet · Mapillary Vistas · traffic sign segmentation · computer vision · deep learning