Pixel-wise gradient boosting treats every pixel as an independent row in a table. That works remarkably well when the signal is spectral — water is dark in near-infrared wherever it is — and it fails when the signal is spatial. A parking lot and a flat roof have nearly identical reflectance; what separates them is shape, size, and what surrounds them. Semantic segmentation models learn that context directly, mapping a multi-band image patch to a per-pixel class map in one forward pass.
This topic is part of Training Geospatial Predictive Models in Python. It picks up where gradient boosting for raster data reaches its ceiling, and it depends on the label rasters produced by rasterizing vector layers for model inputs. The evaluation discipline is the same one demanded everywhere else on this site: without spatial cross-validation strategies, a segmentation score is a measurement of memorisation, not of generalisation.
Problem Framing
A segmentation model is a function from a stack of H × W × B reflectance values to a stack of H × W × C class logits. Three properties of geospatial data make that harder than the natural-image case the architectures were designed for.
Patches are not independent samples. Two 256-pixel patches cut 100 pixels apart share most of their content. Treated as separate rows and split randomly, they leak across the train/validation boundary as thoroughly as duplicated rows would. The split has to happen at a spatial scale much larger than the patch — this is the same phenomenon covered in reducing spatial leakage in model training, just with images instead of points.
Classes are wildly imbalanced and partially unlabelled. A scene may be 80% cropland, 0.4% water, and 30% never digitised at all. Unlabelled is not a class: those pixels must be excluded from the loss, not learned as background. Getting this wrong produces a model that confidently paints “background” over everything a labeller happened to skip.
The model must run on scenes far larger than a patch. Training happens on 256×256 tiles; inference happens on a 10,000×10,000 scene. Naive tiled inference produces a visible grid of seams because a convolution near a patch edge sees zero padding instead of real neighbours. Overlap-and-blend is not an optimisation, it is a correctness requirement.
Prerequisites & Environment Setup
# Pinned requirements for satellite segmentation
torch==2.3.1
torchvision==0.18.1
rasterio==1.3.10
numpy==1.26.4
geopandas==0.14.4
scikit-learn==1.5.0
tqdm==4.66.4
Install with (CPU wheels shown; swap the index URL for a CUDA build):
pip install "torch==2.3.1" "torchvision==0.18.1" --index-url https://download.pytorch.org/whl/cpu
pip install "rasterio==1.3.10" "numpy==1.26.4" "geopandas==0.14.4" \
"scikit-learn==1.5.0" "tqdm==4.66.4"GDAL/PROJ system dependencies (Ubuntu/Debian):
sudo apt-get install -y gdal-bin libgdal-dev libproj-devYou need a feature raster and a label raster that share a grid exactly — same CRS, transform, width and height. Produce them with the workflow in aligning rasterized labels to an existing feature grid and assert the alignment before you write a single line of model code.
Step-by-Step Implementation
Step 1 — Assign spatial blocks, then cut patches inside them
Blocks first, patches second. This ordering is what keeps validation honest.
import numpy as np
import rasterio
def block_assignments(height: int, width: int, block_px: int = 2048,
fractions=(0.7, 0.15, 0.15), seed: int = 0) -> dict:
"""Assign each spatial block to train/val/test. Whole blocks, never patches."""
rows = int(np.ceil(height / block_px))
cols = int(np.ceil(width / block_px))
ids = np.arange(rows * cols)
rng = np.random.default_rng(seed)
rng.shuffle(ids)
n_train = int(round(fractions[0] * len(ids)))
n_val = int(round(fractions[1] * len(ids)))
return {
"train": set(ids[:n_train].tolist()),
"val": set(ids[n_train:n_train + n_val].tolist()),
"test": set(ids[n_train + n_val:].tolist()),
"shape": (rows, cols),
"block_px": block_px,
}
def patch_index(height, width, assign, patch=256, stride=192, split="train"):
"""Yield (row, col) patch origins whose entire footprint lies in one split's blocks."""
rows, cols = assign["shape"]
b = assign["block_px"]
out = []
for r in range(0, height - patch + 1, stride):
for c in range(0, width - patch + 1, stride):
br0, br1 = r // b, (r + patch - 1) // b
bc0, bc1 = c // b, (c + patch - 1) // b
if br0 != br1 or bc0 != bc1:
continue # straddles a block edge — drop it
if (br0 * cols + bc0) in assign[split]:
out.append((r, c))
return outDropping the straddling patches costs a few percent of the training data and buys a validation score you can act on.
Step 2 — Normalise with training-set statistics only
def band_stats(raster_path: str, origins, patch: int = 256, bands=(1, 2, 3, 4, 5, 6)):
"""Per-band mean and std over TRAINING patches only, streamed."""
n = 0
total = np.zeros(len(bands), dtype="float64")
total_sq = np.zeros(len(bands), dtype="float64")
with rasterio.open(raster_path) as src:
for r, c in origins:
win = rasterio.windows.Window(c, r, patch, patch)
arr = src.read(list(bands), window=win).astype("float64")
arr = arr.reshape(len(bands), -1)
total += arr.sum(axis=1)
total_sq += (arr ** 2).sum(axis=1)
n += arr.shape[1]
mean = total / n
std = np.sqrt(np.maximum(total_sq / n - mean ** 2, 1e-12))
return mean.astype("float32"), std.astype("float32")
mean, std = band_stats("features.tif", train_origins)
np.savez("norm_stats.npz", mean=mean, std=std) # ship this with the model
assert (std > 0).all(), "a band has zero variance in the training split"Saving the statistics as a file is what prevents the most common serving bug in this whole workflow — recomputing normalisation from the inference scene, which shifts every input the model sees. The same discipline is described in scaling features consistently between training and inference.
Step 3 — Define the network
import torch
import torch.nn as nn
def conv_block(cin: int, cout: int) -> nn.Sequential:
return nn.Sequential(
nn.Conv2d(cin, cout, 3, padding=1, bias=False),
nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
nn.Conv2d(cout, cout, 3, padding=1, bias=False),
nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
)
class UNet(nn.Module):
"""Compact U-Net sized for multi-spectral patches."""
def __init__(self, in_bands: int = 6, n_classes: int = 5, width: int = 32):
super().__init__()
w = width
self.enc1, self.enc2 = conv_block(in_bands, w), conv_block(w, w * 2)
self.enc3, self.enc4 = conv_block(w * 2, w * 4), conv_block(w * 4, w * 8)
self.pool = nn.MaxPool2d(2)
self.bottleneck = conv_block(w * 8, w * 16)
self.up4 = nn.ConvTranspose2d(w * 16, w * 8, 2, stride=2)
self.dec4 = conv_block(w * 16, w * 8)
self.up3 = nn.ConvTranspose2d(w * 8, w * 4, 2, stride=2)
self.dec3 = conv_block(w * 8, w * 4)
self.up2 = nn.ConvTranspose2d(w * 4, w * 2, 2, stride=2)
self.dec2 = conv_block(w * 4, w * 2)
self.up1 = nn.ConvTranspose2d(w * 2, w, 2, stride=2)
self.dec1 = conv_block(w * 2, w)
self.head = nn.Conv2d(w, n_classes, 1)
def forward(self, x):
e1 = self.enc1(x)
e2 = self.enc2(self.pool(e1))
e3 = self.enc3(self.pool(e2))
e4 = self.enc4(self.pool(e3))
b = self.bottleneck(self.pool(e4))
d4 = self.dec4(torch.cat([self.up4(b), e4], dim=1))
d3 = self.dec3(torch.cat([self.up3(d4), e3], dim=1))
d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
return self.head(d1)
model = UNet(in_bands=6, n_classes=5)
probe = model(torch.zeros(2, 6, 256, 256))
assert probe.shape == (2, 5, 256, 256), f"unexpected output shape {tuple(probe.shape)}"Step 4 — Loss that respects imbalance and ignores unlabelled pixels
IGNORE_INDEX = 255
class DiceCELoss(nn.Module):
"""Class-weighted cross-entropy plus a soft Dice term over labelled pixels only."""
def __init__(self, class_weights: torch.Tensor, dice_weight: float = 0.5):
super().__init__()
self.ce = nn.CrossEntropyLoss(weight=class_weights, ignore_index=IGNORE_INDEX)
self.dice_weight = dice_weight
def forward(self, logits, target):
loss = self.ce(logits, target)
valid = target != IGNORE_INDEX
if valid.any():
probs = logits.softmax(dim=1)
n_classes = logits.shape[1]
safe = target.clone()
safe[~valid] = 0
onehot = torch.nn.functional.one_hot(safe, n_classes).permute(0, 3, 1, 2).float()
m = valid.unsqueeze(1).float()
inter = (probs * onehot * m).sum(dim=(0, 2, 3))
denom = (probs * m).sum(dim=(0, 2, 3)) + (onehot * m).sum(dim=(0, 2, 3))
loss = loss + self.dice_weight * (1 - (2 * inter + 1e-6) / (denom + 1e-6)).mean()
return loss
# Inverse-frequency weights, capped so a 0.1% class does not dominate the gradient.
counts = np.array([4_100_000, 820_000, 96_000, 31_000, 210_000], dtype="float64")
weights = np.clip((counts.sum() / (len(counts) * counts)), 0.2, 8.0)
criterion = DiceCELoss(torch.tensor(weights, dtype=torch.float32))Step 5 — Overlapping-window inference over the full scene
def cosine_window(size: int) -> np.ndarray:
"""2-D Hann taper: full weight at the patch centre, zero at the edge."""
w = np.hanning(size + 2)[1:-1].astype("float32")
return np.outer(w, w)
@torch.no_grad()
def predict_scene(model, raster_path, out_path, mean, std,
patch=256, overlap=64, n_classes=5, device="cpu"):
model.eval().to(device)
taper = cosine_window(patch)
stride = patch - overlap
with rasterio.open(raster_path) as src:
acc = np.zeros((n_classes, src.height, src.width), dtype="float32")
wsum = np.zeros((src.height, src.width), dtype="float32")
for r in range(0, src.height - patch + 1, stride):
for c in range(0, src.width - patch + 1, stride):
win = rasterio.windows.Window(c, r, patch, patch)
x = src.read(window=win).astype("float32")
x = (x - mean[:, None, None]) / std[:, None, None]
logits = model(torch.from_numpy(x)[None].to(device))
probs = logits.softmax(dim=1)[0].cpu().numpy()
acc[:, r:r + patch, c:c + patch] += probs * taper
wsum[r:r + patch, c:c + patch] += taper
profile = src.profile.copy()
labels = np.where(wsum > 0, acc.argmax(axis=0), IGNORE_INDEX).astype("uint8")
profile.update(count=1, dtype="uint8", nodata=IGNORE_INDEX, compress="deflate")
with rasterio.open(out_path, "w", **profile) as dst:
dst.write(labels, 1)
return labelsVerification & Testing
Segmentation demands class-aware metrics. Overall accuracy will tell you the model is excellent right up to the moment somebody looks at the map.
def confusion(pred, target, n_classes, ignore=IGNORE_INDEX):
valid = target != ignore
k = target[valid].astype("int64") * n_classes + pred[valid].astype("int64")
return np.bincount(k, minlength=n_classes ** 2).reshape(n_classes, n_classes)
def iou_per_class(cm):
tp = np.diag(cm).astype("float64")
denom = cm.sum(axis=1) + cm.sum(axis=0) - tp
return np.where(denom > 0, tp / np.maximum(denom, 1), np.nan)
cm = confusion(pred_test, y_test, n_classes=5)
ious = iou_per_class(cm)
print({i: round(float(v), 3) for i, v in enumerate(ious)})
print("macro IoU:", round(float(np.nanmean(ious)), 3))
assert np.nanmin(ious) > 0.0, "a class was never predicted correctly — check class weights"Two structural tests are worth keeping in the repository permanently:
def test_no_patch_crosses_a_split():
"""Every patch footprint must fall inside exactly one block."""
assign = block_assignments(8192, 8192, block_px=2048, seed=0)
for split in ("train", "val", "test"):
for r, c in patch_index(8192, 8192, assign, split=split):
assert (r // 2048) == ((r + 255) // 2048)
assert (c // 2048) == ((c + 255) // 2048)
def test_ignore_pixels_produce_no_gradient():
"""A batch that is entirely unlabelled must give a zero-gradient cross-entropy term."""
logits = torch.randn(1, 5, 8, 8, requires_grad=True)
target = torch.full((1, 8, 8), IGNORE_INDEX, dtype=torch.long)
loss = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)(logits, target)
assert torch.isnan(loss) or float(loss) == 0.0Troubleshooting & Common Errors
CUDA out of memory at batch size 8 — activation memory scales with patch² × width, not with parameter count. Halve the patch size before halving the batch, use torch.cuda.amp.autocast() for mixed precision, and set torch.backends.cudnn.benchmark = True for fixed-size inputs.
Loss becomes nan in the first few hundred steps — nearly always an unmasked nodata sentinel. A -9999 reflectance value survives normalisation as a large negative number and blows up through the batch norms. Convert nodata to NaN at read time and drop or impute those patches.
Validation IoU is much better than test IoU — the validation blocks are adjacent to training blocks. Add a buffer of one block between splits, or use the leave-one-region-out pattern from leave-one-region-out cross-validation with scikit-learn.
Predicted map shows a grid of faint rectangles — tiled inference without overlap and blending. Increase overlap to at least 25% of the patch and weight by the Hann taper shown above.
One class is never predicted — its weight is too low, or it is genuinely too rare for the patch sampler to reach. Over-sample patches containing the rare class rather than raising its weight past about 8, which destabilises training. The trade-offs are the same as in handling class imbalance in land cover classification.
Predictions are shifted by a few pixels against the imagery — the label raster and feature raster transforms differ. Assert transform equality before training; see the alignment guide linked above.
Performance Optimisation
Cut patches once, to disk. Reading 256×256 windows from a compressed GeoTIFF for every epoch makes the data loader the bottleneck. Materialise the training patches as a single memory-mapped .npy stack (or a sharded WebDataset) and the same GPU will run several times faster.
Match num_workers to decode cost, not to core count. With pre-cut patches, two or three workers saturate the GPU; with on-the-fly GeoTIFF reads you may need eight or more, and GDAL_CACHEMAX becomes the limiting factor.
Use mixed precision. torch.autocast roughly halves activation memory and gives a 1.5–2× speedup on any recent GPU, which usually translates into a larger patch — worth more than the raw speed, because a larger patch means more context per prediction.
Export for serving. A trained U-Net is a fixed graph and belongs in ONNX export for geospatial model inference rather than shipping a PyTorch runtime into a GDAL container. Fix the batch dimension to 1 and mark height and width dynamic so a single artifact serves any tile size.
Parallelise inference by tile, not by pixel. Full-scene prediction is embarrassingly parallel over tiles once you keep the overlap bookkeeping straight — the distribution pattern in scaling batch inference with Dask and Ray applies directly.
FAQ
Why does my model score 95% accuracy but produce useless maps?
Overall pixel accuracy is dominated by the majority class: if 93% of pixels are background, predicting background everywhere scores 93%. Report per-class intersection over union and the macro mean, and use the class-weighted plus Dice loss above so rare classes generate real gradients.
Can I use random patch splits instead of spatial blocks?
Not if you want a number you can trust. Adjacent patches share pixels and context, so random splitting places near-duplicates on both sides of the boundary and typically inflates IoU by 20–30 points. Split blocks first, cut patches second, and drop the patches that straddle a boundary.
How do I avoid visible seams in a full-scene prediction?
Predict on overlapping tiles and blend with a Hann taper, as in Step 5. Convolutions near a patch edge see padding rather than real neighbours, so edge predictions are systematically worse; the taper down-weights exactly those pixels.
How many labelled patches do I need?
Fewer than intuition suggests, because each patch contributes tens of thousands of labelled pixels. A few hundred well-distributed 256×256 patches per class often beats a tuned pixel-wise baseline — provided they span the terrain, season and sensor conditions production will see.
Related
- Building a U-Net for Land Cover Segmentation in PyTorch — the architecture and training loop in full
- Tiling Large Rasters into Training Patches — patch extraction that does not leak
- Building a CNN for Satellite Imagery Classification — the patch-classification counterpart
- Rasterizing Vector Layers for Model Inputs — producing the label raster this workflow consumes
- Spatial Cross-Validation Strategies — the evaluation discipline behind the block split