Introduction to Geospatial AI¶
2025-11-15 — Workshop conducted at NIE for geospatial AI performing instance segmentation on high res imagery
Geospatial Instance Segmentation with TorchGeo + Mask R-CNN¶
NWPU VHR-10 (Very High Resolution, 10 classes)¶
Link to github repository for installation steps on local here
Goal for this workshop
By the end, you'll be able to:
- Load the NWPU VHR-10 dataset from TorchGeo (remote sensing, 10 object classes).
- Explore images, bounding boxes, and instance masks.
- Build a Mask R-CNN model (from torchvision) for instance segmentation.
- Train on a small subset (for speed) and visualize predictions.
We'll use:
- TorchGeo for geospatial dataset handling.
- torchvision Mask R-CNN for instance segmentation.
0. Environment setup¶
Run this once at the top of your notebook.
⚠️ Warning for instructors:
- VHR-10 needs the
datasetsextra, pluspycocotoolsfor COCO-style annotations. - Always request a GPU runtime in Colab (T4 is enough).
💡 Idea: Have students run the pip cell themselves so they see what libraries are used.
In [ ]:
# If you're on Colab, uncomment this cell:
!pip install -q "torchgeo[datasets]" torch torchvision matplotlib
import torch
import torchvision
from torch.utils.data import DataLoader, Subset
import matplotlib.pyplot as plt
from torchgeo.datasets import VHR10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Using device:", device)
# 🧪 Practice 0:
# 1. Run this cell and see if you're on CPU or GPU.
# 2. If you have a GPU: print(torch.cuda.device_count()) to see how many devices are available.
print("CUDA device count:", torch.cuda.device_count())
1. Load and explore the VHR-10 dataset¶
VHR-10 (NWPU VHR-10) is a very high resolution (0.08–2 m) remote sensing dataset with:
- 800 images total
- 650 "positive" images (contain at least one object)
- 150 "negative" images (no objects)
- 10 object classes: airplane, ship, storage tank, baseball diamond, tennis court, basketball court, ground track field, harbor, bridge, vehicle.
TorchGeo provides:
- images as tensors (C, H, W)
- COCO-style annotations converted into
boxes,labels, andmasks.
⚠️ Important:
- Use
split="positive"for detection/segmentation (only positives have annotations).
In [ ]:
root = "data" # change if you want
ds = VHR10(root=root, split="positive", download=True, checksum=False)
print("Dataset length (positive images):", len(ds))
sample = ds[0]
print("Keys in one sample:", sample.keys())
print("Image shape:", sample["image"].shape) # (C, H, W)
print("Boxes shape:", sample["bbox_xyxy"].shape)
print("Labels shape:", sample["label"].shape)
print("Has masks:", "mask" in sample)
Visualize an image with ground truth boxes & masks¶
TorchGeo's plot method is built in for VHR-10 and can display:
- RGB image
- bounding boxes
- masks (if present)
This is perfect to see what instance segmentation labels look like.
In [ ]:
ds.plot(sample) # draws image with GT boxes + masks
plt.show()
# 🧪 Practice 1:
# 1. Pick a random index, e.g. idx = 42, and visualize that sample with ds.plot(ds[idx]).
# 2. Print the number of objects in that image: len(sample["labels"]).
# Example (uncomment):
# idx = 42
# sample = ds[idx]
# print("Number of objects:", len(sample["labels"]))
# ds.plot(sample); plt.show()
2. Simple preprocessing & DataLoader¶
Mask R-CNN in torchvision expects:
- images as
float32tensors in [0, 1] range - a list of images (variable size allowed)
- a list of target dicts, each with:
boxes:[N, 4]labels:[N]masks(optional):[N, H, W]
TorchGeo already:
- returns
imageas float tensor (0–255) - gives
boxes,labels,masksper sample
We’ll:
- Normalize
image / 255.0using a small transform function. - Use a custom
collate_fnthat returns a list of samples (no stacking), since image sizes differ.
In [ ]:
def transform_to_unit_range(sample: dict) -> dict:
# Normalize image to [0, 1] range.
image = sample["image"] / 255.0
sample["image"] = image
return sample
# Recreate dataset with transform
ds = VHR10(root=root, split="positive", download=False, transforms=transform_to_unit_range)
sample = ds[0]
print("Image min/max after normalization:",
sample["image"].min().item(), sample["image"].max().item())
def collate_fn(batch):
# Simple collate function for detection / instance segmentation.
# Input: list of samples (dicts)
# Output: same list (we handle conversion in the training loop).
return batch
# For speed in workshop, use a small subset (e.g. first 100 images)
subset_indices = list(range(100))
train_dataset = Subset(ds, subset_indices)
train_loader = DataLoader(
train_dataset,
batch_size=2,
shuffle=True,
num_workers=2,
collate_fn=collate_fn,
)
print("Batches in train_loader:", len(train_loader))
# 🧪 Practice 2:
# Inspect one batch and print image shapes:
# batch = next(iter(train_loader))
# for i, s in enumerate(batch):
# print(i, s["image"].shape)
3. Build Mask R-CNN (torchvision)¶
We’ll use maskrcnn_resnet50_fpn from torchvision:
- Backbone: ResNet-50 + FPN
- Detection head: region proposals + class & box prediction
- Segmentation head: mask prediction per instance
VHR-10 has 10 classes + background, and VHR10.categories includes 'background' at index 0.
In [ ]:
from torchvision.models.detection import (
maskrcnn_resnet50_fpn,
MaskRCNN_ResNet50_FPN_Weights,
)
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
# Get class count from dataset (handle Subset)
if isinstance(train_dataset, Subset):
base_ds = train_dataset.dataset # underlying VHR10
else:
base_ds = train_dataset
num_classes = len(base_ds.categories)
print("Number of classes (including background):", num_classes)
# 1) Start from COCO-pretrained weights (good initialization)
weights = MaskRCNN_ResNet50_FPN_Weights.COCO_V1
model = maskrcnn_resnet50_fpn(weights=weights)
model.to(device)
# 2) Replace classification head (box predictor)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# 3) Replace mask head
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
model.roi_heads.mask_predictor = MaskRCNNPredictor(
in_features_mask,
hidden_layer,
num_classes,
)
model.to(device)
print("Model ready.")
# ⚠️ Note: We keep most COCO-trained weights and only reinitialize the last heads.
4. Training loop (mini version for workshop)¶
For each batch:
- Convert list of samples →
images(list of tensors) +targets(list of dicts). - Forward pass:
loss_dict = model(images, targets). - Sum loss terms and backprop.
We’ll only train for 1–2 epochs on a subset of 100 images so it runs quickly.
In [ ]:
# Optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005)
num_epochs = 2
for epoch in range(num_epochs):
model.train()
epoch_loss = 0.0
for batch_idx, batch in enumerate(train_loader):
# batch is a list of samples (dicts)
images = [b["image"].to(device) for b in batch]
targets = []
for b in batch:
target = {
"boxes": b["bbox_xyxy"].to(device),
"labels": b["label"].to(device),
}
if "mask" in b:
target["masks"] = b["mask"].to(device)
targets.append(target)
loss_dict = model(images, targets)
loss = sum(loss_dict.values())
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.item()
if (batch_idx + 1) % 10 == 0:
print(
f"Epoch [{epoch+1}/{num_epochs}] "
f"Batch [{batch_idx+1}/{len(train_loader)}] "
f"Loss: {loss.item():.4f}"
)
avg_loss = epoch_loss / len(train_loader)
print(f"Epoch [{epoch+1}/{num_epochs}] Average loss: {avg_loss:.4f}")
# 🧪 Practice 3:
# - Change num_epochs to 1 or 3 and observe training time and loss.
# - Change batch_size in the DataLoader to 4 and see if GPU memory is okay.
5. Inference & visualization¶
Now we’ll:
- Take a few samples from the dataset.
- Run
modelineval()mode without targets → predictions. - Plot:
- image
- ground-truth boxes
- predicted boxes + masks + scores
We’ll reuse TorchGeo’s plot method and feed predictions back into the sample dict using keys:
prediction_boxesprediction_labelsprediction_scoresprediction_masks
In [ ]:
sample
Out[ ]:
In [ ]:
sample_for_plot
Out[ ]:
In [ ]:
In [ ]:
model.eval()
indices_to_show = [0, 10, 20] # change as you like
for idx in indices_to_show:
sample = ds[idx]
image = sample["image"].to(device)
with torch.no_grad():
prediction_list = model([image])
prediction = prediction_list[0]
# Prepare dict for plotting
sample_for_plot = {
"image": sample["image"].cpu(),
"bbox_xyxy": sample["bbox_xyxy"].cpu(),
"label": sample["label"].cpu(),
}
if "mask" in sample:
sample_for_plot["masks"] = sample["mask"].cpu()
# Copy predictions
sample_for_plot["prediction_boxes"] = prediction["boxes"].cpu()
sample_for_plot["prediction_labels"] = prediction["labels"].cpu()
sample_for_plot["prediction_scores"] = prediction["scores"].cpu()
if "masks" in prediction:
# prediction["masks"]: [N, 1, H, W] → [N, H, W]
sample_for_plot["prediction_masks"] = prediction["masks"].cpu().squeeze(1)
fig = ds.plot(sample_for_plot, show_feats="both")
plt.suptitle(f"Sample {idx} — GT vs. Mask R-CNN predictions")
plt.show()
# 🧪 Practice 4:
# Add a confidence threshold, e.g. keep only scores >= 0.5:
#
# scores = prediction["scores"]
# keep = scores >= 0.5
# sample_for_plot["prediction_boxes"] = prediction["boxes"][keep].cpu()
# sample_for_plot["prediction_labels"] = prediction["labels"][keep].cpu()
# sample_for_plot["prediction_scores"] = scores[keep].cpu()
# sample_for_plot["prediction_masks"] = prediction["masks"][keep].cpu().squeeze(1)
#
# Try different thresholds (0.3, 0.7, 0.9) and see how the visualization changes.
6. Simple evaluation: box IoU, precision, recall¶
We’ll compute a very small evaluation:
- Use a tiny validation set (e.g. 20 images).
- For each image:
- Run the model in
eval()mode. - Filter predictions by score (e.g. ≥ 0.5).
- Match predicted boxes to ground-truth boxes using IoU ≥ 0.5 and same class.
- Run the model in
- Count:
- TP (true positives): matched predictions
- FP (false positives): predictions that didn’t match any GT
- FN (false negatives): GT objects with no matching prediction
Then compute:
- Precision = TP / (TP + FP)
- Recall = TP / (TP + FN)
This is not full COCO-style mAP, but it’s easy to understand and enough for a workshop.
In [ ]:
from torchvision.ops import box_iou
# Build a tiny validation set: next 20 images after the training subset
val_indices = list(range(100, min(120, len(ds))))
val_dataset = Subset(ds, val_indices)
val_loader = DataLoader(
val_dataset,
batch_size=1, # easier for teaching + matching logic
shuffle=False,
num_workers=2,
collate_fn=collate_fn,
)
def evaluate_model(model, data_loader, score_thresh=0.5, iou_thresh=0.5):
model.eval()
total_tp = 0
total_fp = 0
total_fn = 0
with torch.no_grad():
for batch in data_loader:
# batch is a list of size 1 because batch_size=1
sample = batch[0]
image = sample["image"].to(device)
gt_boxes = sample["boxes"].to(device)
gt_labels = sample["labels"].to(device)
# Forward pass (no targets → predictions)
outputs = model([image])[0]
pred_boxes = outputs["boxes"].to(device)
pred_labels = outputs["labels"].to(device)
scores = outputs["scores"].to(device)
# 1) Filter predictions by score
keep = scores >= score_thresh
pred_boxes = pred_boxes[keep]
pred_labels = pred_labels[keep]
if pred_boxes.numel() == 0:
# no predictions at this threshold
total_fn += gt_boxes.shape[0]
continue
# 2) Compute IoU between all GT and predicted boxes
ious = box_iou(gt_boxes, pred_boxes) # [num_gt, num_pred]
# 3) Greedy matching based on max IoU per GT
matched_pred = set()
tp = 0
for gt_idx in range(gt_boxes.shape[0]):
# Best predicted box for this GT
iou_row = ious[gt_idx]
best_iou, best_pred_idx = iou_row.max(0)
# Check if it’s a good match (IoU + same label)
if (
best_iou >= iou_thresh
and best_pred_idx.item() not in matched_pred
and pred_labels[best_pred_idx] == gt_labels[gt_idx]
):
tp += 1
matched_pred.add(best_pred_idx.item())
fp = pred_boxes.shape[0] - len(matched_pred)
fn = gt_boxes.shape[0] - tp
total_tp += tp
total_fp += fp
total_fn += fn
precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0
recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
return {
"tp": total_tp,
"fp": total_fp,
"fn": total_fn,
"precision": precision,
"recall": recall,
}
metrics = evaluate_model(model, val_loader, score_thresh=0.5, iou_thresh=0.5)
print("Evaluation on small validation set:")
print(metrics)
In [ ]:
# 🧪 Practice 5:
# 1. Try different score thresholds: 0.3, 0.5, 0.7.
# metrics_03 = evaluate_model(model, val_loader, score_thresh=0.3)
# metrics_07 = evaluate_model(model, val_loader, score_thresh=0.7)
# 2. Compare precision and recall:
# - What happens to precision when you increase the threshold?
# - What happens to recall?
#
# 3. (Discussion) How would you choose a threshold for a *real* application?
6. Wrap-up & extensions¶
Key takeaways:
- Geospatial datasets can be handled with TorchGeo, which provides:
- dataset classes
- metadata
- plotting helpers
- Instance segmentation = object detection + per-instance mask (Mask R-CNN).
- Mask R-CNN from torchvision works on remote sensing data with:
- proper image scaling to
[0, 1] - correct
num_classes - a custom
collate_fnfor variable-size images.
- proper image scaling to
Extensions:
- Use TorchGeo’s
InstanceSegmentationTask+ PyTorch Lightning for cleaner training and metrics. - Add data augmentations (random flip, resize, color jitter) and study their effect.
- Compare with:
- object detection only (Faster R-CNN)
- semantic segmentation models (no instance separation)
- YOLO-style models trained on VHR-10.