Module 7: Segmentation Networks for Imagery

Phase: III, Deep learning
Level: Advanced
Estimated time: ~12 hours (about 4 h concepts, 7 h lab, 1 h write-up)
Prerequisites: Phase II (Modules 4 to 6), Phase I, Phase 0.
Portfolio thread: adds a learned model to Portfolio Project A, a trained U-Net plus a predict.py inference CLI, alongside the index baseline from Module 2.

Why this module

U-Net on HLS burn scars is the canonical geospatial fine-tuning task; it is exactly what NASA-IMPACT used to demonstrate the Prithvi foundation model. Build this baseline yourself first, before touching a foundation model in Module 8, so that when you later compare them you can say with evidence whether the foundation model actually helped or just cost more. Segmentation is also the workhorse of pixel-level labeling across the field (burn scars, water, deforestation, crops), so a clean, tested segmentation pipeline is a reusable asset and a strong portfolio piece.

Learning outcomes

After this module you can:

1. Explain semantic segmentation and the U-Net architecture, and name when to reach for FPN, DeepLabV3+, or UPerNet instead.
2. Choose loss functions for class-imbalanced segmentation (Dice, focal, Lovasz, and combinations).
3. Use TorchGeo and segmentation_models_pytorch, with correct multispectral normalization and augmentation.
4. Train and evaluate a U-Net on the HLS burn-scar dataset, logging IoU and F1.
5. Run tiled sliding-window inference on your own AOI and compare to the dNBR and MTBS baselines.
6. Write a predict.py CLI that goes from a STAC query to a COG, the shape of a production inference job.

Concept lessons (about 4 hours)

Lesson 1: Semantic segmentation and U-Net

Semantic segmentation assigns a class to every pixel. U-Net is the classic architecture: a contracting encoder that captures context at coarser and coarser scales, an expanding decoder that recovers full resolution, and skip connections that carry fine spatial detail from encoder to decoder so edges stay sharp. The encoder is usually a standard backbone (for example a ResNet). Alternatives worth knowing: FPN (feature pyramid, multi-scale), DeepLabV3+ (atrous convolutions for large receptive fields), and UPerNet (often paired with transformer backbones, which matters in Module 8). U-Net remains a strong, fast default for EO segmentation.

Lesson 2: Losses for class imbalance

Burn scars, water, and deforestation are usually a small fraction of the pixels, so plain cross-entropy is dominated by the background and the model learns to predict "nothing." Use losses that focus on the rare class:

• Dice loss directly optimizes overlap and is robust to imbalance.
• Focal loss down-weights easy, confident background pixels so training attends to hard cases.
• Lovasz-softmax optimizes the IoU metric more directly.
• Combinations (Dice plus cross-entropy, or Dice plus focal) are common and strong.

Evaluate with IoU (Jaccard) and F1 (Dice) per class, never overall pixel accuracy, which the background dominates.

Lesson 3: Multispectral deep learning in practice

A few things trip up newcomers coming from natural-image deep learning:

• Do not use ImageNet statistics or three RGB bands. EO imagery has more bands and different value ranges; compute band-wise mean and standard deviation from your data and adapt the network's first convolution to the number of input channels.
• Augment sensibly. For nadir imagery, flips and 90-degree rotations are free and physically valid (no privileged orientation), unlike natural photos. Be careful with augmentations that distort band relationships.
• Sample patches deliberately. Use TorchGeo's datasets, samplers, and datamodules to tile scenes into training patches, and make sure patches from the same scene do not straddle your train and validation splits, or you leak.

Lesson 4: Training and evaluation discipline

Carry the spatial-holdout lesson from Module 4 into deep learning: split by geography, not by random tiles from the same scene, or neighboring patches leak between train and test and your metrics lie. Keep a reproducible config, track runs (Weights and Biases or MLflow), and run a small learning-rate and loss sweep rather than one lucky run. Report IoU and F1 on a genuinely held-out area.

Lesson 5: Inference and productionization

Real scenes are far larger than training patches, so you infer with a sliding window: tile the scene with overlap, predict each tile, and blend the overlaps to avoid visible seams. Then wrap it in a CLI: predict.py takes a STAC query, builds the cube (Module 1), runs the model, and writes a COG. That CLI is the shape of a production inference job and sets up Phase IV, so build it cleanly now.

Guided lab (about 7 hours): a U-Net burn-scar segmenter

Step 1: Data and datamodule

Download the NASA-IMPACT HLS burn-scar dataset (on Hugging Face under ibm-nasa-geospatial) and build a TorchGeo or Lightning datamodule with band-wise normalization and flip/rotation augmentation.

Step 2: Pure helpers with tests

These are the numeric guts you can test without a GPU:

"""seg.py: segmentation helpers (pure, testable)."""
from __future__ import annotations
import numpy as np

def iou_score(pred, target, eps=1e-6):
    """Intersection over union for binary masks."""
    pred = np.asarray(pred, bool); target = np.asarray(target, bool)
    inter = (pred & target).sum()
    union = (pred | target).sum()
    return float((inter + eps) / (union + eps))

def normalize_bands(x, mean, std):
    """Band-wise normalization for a (C, H, W) array (never ImageNet stats)."""
    mean = np.asarray(mean, float)[:, None, None]
    std = np.asarray(std, float)[:, None, None]
    return (np.asarray(x, float) - mean) / std

def tile_indices(H, W, size, stride):
    """Top-left corners for sliding-window tiling that always covers the edges."""
    def axis(n):
        if n <= size:
            return [0]
        idx = list(range(0, n - size + 1, stride))
        if idx[-1] != n - size:
            idx.append(n - size)
        return idx
    return [(y, x) for y in axis(H) for x in axis(W)]

# tests/test_seg.py
import numpy as np
from eo_portfolio.seg import iou_score, normalize_bands, tile_indices

def test_iou():
    a = np.array([[1, 1], [0, 0]]); b = np.array([[1, 0], [0, 0]])
    assert abs(iou_score(a, b) - 0.5) < 1e-3   # inter 1, union 2
    assert iou_score(a, a) > 0.99

def test_normalize_zeroes_at_mean():
    x = np.ones((2, 3, 3))
    assert np.allclose(normalize_bands(x, [1, 1], [2, 2]), 0.0)

def test_tiles_cover_edges():
    t = tile_indices(10, 10, 5, 5)
    assert (0, 0) in t and (5, 5) in t and len(t) == 4
    t2 = tile_indices(10, 10, 6, 6)
    assert (0, 0) in t2 and (4, 4) in t2   # last tile snapped to cover the edge

Step 3: Train the U-Net

Train a U-Net with a ResNet-34 encoder from segmentation_models_pytorch, using a Dice-plus-focal loss and band-wise normalization. Log IoU and F1 per epoch to Weights and Biases or MLflow, and run a small learning-rate and loss sweep. Keep the split geographic.

Step 4: Inference on your AOI

Pull HLS chips over your fire AOI (Module 1), run tiled sliding-window inference using tile_indices, blend overlaps, and write the predicted burn mask as a COG. Compare it to the Module 2 dNBR map and MTBS: where does the learned model beat the index, and where does it not?

Step 5: predict.py CLI and commit

Write predict.py that takes a STAC query and a checkpoint and writes a COG. Commit the config, the metrics table, seg.py, tests, the notebook, and the CLI.

git add src/eo_portfolio/seg.py src/eo_portfolio/predict.py tests/test_seg.py notebooks/09_unet_burnscar.ipynb
git commit -m "Module 7: U-Net burn-scar segmenter, tested helpers, predict CLI"
git push

Checkpoint

Self-check (answers below).

1. What do the skip connections in a U-Net do, and why do they matter for sharp masks?
2. Why is pixel accuracy a poor metric for burn-scar segmentation, and what do you use instead?
3. Why must you not use ImageNet normalization for HLS imagery?
4. Why tile with overlap and blend at inference instead of predicting tile by tile with no overlap?

Interview-style questions (practice out loud).

• Walk through training a segmentation model on an imbalanced EO dataset, from splits to loss to metrics.
• Your validation IoU is high but predictions on a new region are poor. Give the likely causes.
• Why build a U-Net baseline before fine-tuning a geospatial foundation model?
• How would you turn a trained model into a production inference job over arbitrary AOIs?

Answers. (1) Skip connections pass high-resolution encoder features to the decoder so fine spatial detail and edges are preserved that pooling would otherwise lose.
(2) Background dominates, so pixel accuracy is high even for a model that misses the scar; use IoU and F1 on the burn class.
(3) HLS has different bands and value ranges than natural RGB images; ImageNet statistics are wrong, so compute band-wise statistics from your data and adapt the first convolution.
(4) Tile-by-tile prediction produces visible seams at borders where the model has no context; overlapping and blending removes the seams.

Deliverable

A tested seg.py (IoU, band-wise normalization, tiling), a trained U-Net with a reproducible config and an IoU/F1 metrics table, a predict.py CLI that goes from STAC query to COG, and a notebook comparing the learned model to dNBR and MTBS over your AOI. Portfolio Project A now has a learned model beside the index baseline.

What a hiring manager sees

A trained U-Net is common; a disciplined one is not. Geographic splits, imbalance-aware losses, correct multispectral normalization, IoU and F1 reporting, and a real inference CLI show you can build segmentation that survives contact with new regions. Having a strong classical baseline (dNBR) and a learned model side by side, with an honest comparison, is exactly the judgment employers want before anyone spends money on a foundation model.

Currency note

Verified August 2026. TorchGeo is at version 0.9.0 (released 14 February 2026, requiring Python 3.12+ and PyTorch 2.2+) and is now community and OSGeo governed under the torchgeo/torchgeo organization (which also governs TorchGeo-Bench and TerraTorch); there is no 0.10 release yet, so use the current 0.9.x and confirm the datamodule and sampler APIs, which change between versions (0.9 removed the crs output key, made bounds a tensor, added a transform key, and switched point datasets to a keypoints key). TorchGeo 0.9 also added one-click "open in Lightning Studios" buttons on its tutorial notebooks. segmentation_models_pytorch and PyTorch Lightning are stable. The NASA-IMPACT HLS burn-scar dataset is on Hugging Face under ibm-nasa-geospatial; confirm the current dataset card and license. Pin all versions in your Module 0.1 environment file, and note GPU is needed for training (a free-tier notebook GPU is enough for this dataset).