Pixel-level semantic segmentation of aerial/drone imagery into 5 land-cover classes (background, building, woodland, water, road) with a U-Net trained in PyTorch on the LandCover.ai dataset — from a from-scratch baseline to a pretrained-encoder model, evaluated rigorously and shipped as a deployed demo.
Test mIoU: 0.7224 | Pixel accuracy: 0.9217 | +0.19 mIoU over baseline
Left:- input | Center:- segmentation overlay | Right: predicted class map
- Problem
- Research question
- Dataset
- Models
- Tools
- Pipeline
- Architecture
- Results
- Deployment
- Project structure
- Reproducing the model
- Future improvements
- Portfolio summary
- License
- Acknowledgements
pip install -r requirements.txt
python app.py # opens local Gradio UI at http://127.0.0.1:7860Or try it online: Open demo.
Given an aerial/drone image, assign every pixel to one of five land-cover
classes:- background, building, woodland, water, road.
Aerial imagery is heavily class-imbalanced (background/vegetation dominate), so the main challenge is recovering the rare but important classes (building, road). Outputs support urban planning (tracking green space, paved area, building footprint), disaster response (pre/post-event damage extent), and infrastructure monitoring (flagging unauthorized construction).
How much does an ImageNet-pretrained EfficientNet-B0 encoder improve semantic segmentation over a U-Net trained from scratch on the same data?
The project compares a from-scratch U-Net baseline against a U-Net with a pretrained EfficientNet-B0 encoder under identical data, loss, and evaluation conditions.
Dataset: LandCover.ai:- 41 aerial orthophotos of Poland covering 216 km²: 33 at 25 cm resolution (~9000×9500 px) and 8 at 50 cm (~4200×4700 px). Tiled into 512×512 patches with the official splitter (stride-512 grid, edge tiles dropped), so results are comparable to published benchmarks.
| Split | Tiles |
|---|---|
| Train | 7,470 |
| Validation | 1,602 |
| Test | 1,602 |
| Total | 10,674 |
Classes (single-channel masks, pixel value = class id), with area coverage:-
| Class | Area (km²) | Label |
|---|---|---|
| Background | 125.75 | 0 |
| Building | 1.85 | 1 |
| Woodland | 72.02 | 2 |
| Water | 13.15 | 3 |
| Road | 3.50 | 4 |
Background + woodland cover ~92% of pixels while building (~0.9%) and road (~1.6%) are rare — this imbalance motivates the weighted losses used later.
| Baseline | Improved | |
|---|---|---|
| Architecture | U-Net from scratch | U-Net + EfficientNet-B0 (ImageNet-pretrained, via segmentation-models-pytorch) |
| Loss | plain cross-entropy | 0.5·Dice + 0.5·Focal (γ=2) |
| Extras | — | LR / augmentation sweeps |
| Test mIoU | 0.5316 | 0.7224 |
| Tool | Used for |
|---|---|
PyTorch, segmentation-models-pytorch |
model definition, training, TorchScript export |
| Albumentations | image/mask augmentation pipeline |
| U-Net (from scratch) → EfficientNet-B0 | baseline vs pretrained backbone |
| Dice + Focal loss | class-imbalance handling |
| Gradio | upload-and-infer demo |
| Kaggle | week 1–2 notebooks (EDA, data engineering) |
| molab (marimo) | week 3–5 notebooks (training, evaluation) |
| Hugging Face Spaces | deployed demo |
| # | Stage | What happens |
|---|---|---|
| 1 | Data prep | Tile 41 orthophotos into 10,674 512² patches (official splitter), 70/15/15 split |
| 2 | Augmentation | Flips / rotation / color jitter via Albumentations, applied identically to image + mask |
| 3 | Baseline | U-Net from scratch + cross-entropy → test mIoU 0.53 |
| 4 | Improvement | EfficientNet-B0 encoder + Dice/Focal loss, LR/aug sweeps → test mIoU 0.72 |
| 5 | Evaluation | Per-class metrics, confusion matrix, robustness checks |
| 6 | Deploy | Export to TorchScript → Gradio demo → Hugging Face Spaces |
U-Net from scratch with encoder channels: 32 → 64 → 128 → 256 → 512, skip connections, and transposed-convolution decoder.
U-Net with EfficientNet-B0 encoder (ImageNet-pretrained); decoder stages up-convolve and concat the matching encoder stages via skip connections.
Input images larger than 512 px are tiled with 64 px overlap; logits are averaged across overlapping tiles before argmax.
Held-out test set (1,602 tiles, no overlap with training):-
| class | IoU | Dice | precision | recall |
|---|---|---|---|---|
| background | 0.8751 | 0.9334 | 0.9759 | 0.8945 |
| building | 0.5203 | 0.6845 | 0.5454 | 0.9188 |
| woodland | 0.8663 | 0.9284 | 0.8954 | 0.9638 |
| water | 0.8901 | 0.9419 | 0.9197 | 0.9651 |
| road | 0.4600 | 0.6301 | 0.5032 | 0.8428 |
| mean | 0.7224 | — | acc 0.9217 | — |
The pretrained model improved mean IoU by +0.19 (0.5316 -> 0.7224) over the from-scratch baseline. The remaining weak classes are building and road — small-object classes with high recall but lower precision, the classic imbalance/boundary failure mode.
Every run shares the same tiles, splits, and test set (logged in
experiment_log.csv), so each step's
contribution is isolated:
| Run | Encoder | Loss | Test mIoU |
|---|---|---|---|
| W3 baseline | U-Net from scratch | cross-entropy | 0.5316 |
| A | EfficientNet-B0 | weighted cross-entropy | 0.6720 |
| B | EfficientNet-B0 | 0.5·Dice + 0.5·Focal | 0.7046 |
| Final | EfficientNet-B0 | Dice+Focal + LR/aug sweeps | 0.7224 |
Pretraining alone (A) adds +0.14; Dice+Focal on top (B) adds another +0.03, mostly by rescuing building (0.42 → 0.51) and road (0.36 → 0.41) IoU.
Best 10 predictions (left) and worst 10 (right) on the test set:-
Performance under brightness/contrast shift, Gaussian noise, and combined perturbations:-
Full evaluation details (failure analysis, robustness data): WEEK5_EVALUATION.md.
The best model was exported to TorchScript and verified against the original PyTorch model:
models/
└── traced_unet.pt (25 MB, CPU inference ~1 s/tile)
It runs behind a Gradio app: upload an aerial image → colorized 5-class overlay.
Live demo:- Open demo
pip install -r requirements.txt
python app.py # opens local Gradio UI at http://127.0.0.1:7860├ app.py # Gradio demo (TorchScript inference)
├ export_model.py # .pth -> traced TorchScript module
├ utils.py # tiling, dataset, training utilities
├ requirements.txt # minimal demo deps
├ WEEK5_EVALUATION.md # metrics, failure modes, robustness
├ models/
│ └── traced_unet.pt # exported model
├ assets/ # example outputs (demo, galleries, figures)
├ notebooks/ # week-1 EDA … week-5 evaluation
└ aerial-segmentation-efficientnet-b0/ # HF Space (Gradio, ZeroGPU)
Training config (best run, notebooks/week-4-finetuning.ipynb):-
| Setting | Value |
|---|---|
| Input size | 512 × 512 (RGB, uint8, /255) |
| Batch size | 128 |
| Optimizer | Adam, lr 5e-5 (sweep also tried 1e-4) |
| Scheduler | ReduceLROnPlateau on val mIoU, patience 3 |
| Epochs | 30 (early stopping, patience 6) |
| Loss | 0.5 · Dice + 0.5 · Focal (γ=2), class-weighted |
| Augmentation | RandomBrightnessContrast, RandomRotate90, HorizontalFlip (p=0.5) |
| Best val / test mIoU | 0.7124 / 0.7224 |
# 1. tiles + splits: utils.py (split_images) from notebooks/week-3-unet.ipynb
# 2. train: notebooks/week-4-finetuning.ipynb (writes week4_outputs/)
python export_model.py # best_week4_model.pth -> models/traced_unet.ptThe traced module expects the exact training preprocessing: uint8 BGR, / 255,
no ImageNet normalization (app.py replicates this).
Week-by-week journey:
| Week | Notebook | What was done |
|---|---|---|
| 1 | week-1-sample-eda.ipynb |
EDA on the sample dataset |
| 2 | week-2-data-engineering-landcoverai.ipynb |
Tiling, augmentation, dataloaders |
| 3 | week-3-unet.ipynb |
U-Net baseline |
| 4 | week-4-finetuning.ipynb |
Pretrained encoder, Dice+Focal, sweeps |
| 5 | week-5-evaluation.ipynb |
Full evaluation & analysis |
- Tiled inference for arbitrarily large aerial images (already done in the demo; extend to evaluation)
- Stronger class-imbalance handling (class-weighted sampling, oversampling rare-class tiles)
- Boundary-aware losses (e.g. boundary-weighted Dice) for road/building edges
- Additional backbones (ResNet34, DeepLabV3+ head) and ensembling
- Multi-class thresholding / post-processing to close the road precision gap
A U-Net with an ImageNet-pretrained EfficientNet-B0 encoder classifying every pixel of aerial imagery into 5 land-cover classes at 0.72 mean IoU (up from a 0.53 from-scratch baseline), built end-to-end — data pipeline, training, evaluation — and shipped as a TorchScript model behind a deployed Gradio demo.
MIT — see LICENSE. LandCover.ai is research-use.
Implementation and learnings based on this notebook, which provides a step-by-step implementation with explanation.