# Copyright (c) TorchGeo Contributors. All rights reserved.
# Licensed under the MIT License.

NAIP Road Segmentation

Written by: Isaac Corley

Introduction

The objective of this tutorial is to run inference with TorchGeo’s pre-trained road segmentation model on NAIP imagery.

This walkthrough covers:

• How to stream a sample NAIP tile and road mask directly from remote GeoTIFF assets;

• How to load Unet_Weights.NAIP_RGBN_RESNET18_CHESAPEAKERSC and its preprocessing pipeline;

• How to run forward inference and compare predicted road masks against ground truth.

Unlike a training tutorial, this notebook focuses only on model loading, preprocessing, and inference.

References

Seeing the roads through the trees: A benchmark for modeling spatial dependencies with aerial imagery, Robinson et al (2024)

ChesapeakeRSC dataset and implementation resources: Hugging Face dataset card | GitHub source code

Environment

For the environment, install TorchGeo. Matplotlib is already available as a core dependency in the TorchGeo stack used by this tutorial.

!uv pip install torchgeo

Imports

We import:

• rasterio to read remote GeoTIFFs;

• torch and numpy for tensor preparation;

• matplotlib for visualization;

• TorchGeo model weights + model constructor for inference.

import matplotlib.pyplot as plt
import numpy as np
import rasterio
import torch

from torchgeo.models import Unet_Weights, unet

Dataset

We use one NAIP image tile and one corresponding road mask from the ChesapeakeRSC assets.

To keep the tutorial lightweight, both files are read directly from remote HTTPS URLs in rasterio (no local download step).

image_url = 'https://github.com/isaaccorley/ChesapeakeRSC/raw/66b29f478ee5263f51a8620eb3a34093a970523d/assets/8782_image.tif'
mask_url = 'https://github.com/isaaccorley/ChesapeakeRSC/raw/66b29f478ee5263f51a8620eb3a34093a970523d/assets/8782_mask.tif'

with rasterio.open(image_url) as src:
    image_np = src.read().astype(np.float32)
image = torch.from_numpy(image_np).unsqueeze(0)

with rasterio.open(mask_url) as src:
    mask = torch.from_numpy(src.read(1))
    target = (mask == 9) | (mask == 12)  # road=9, road-occluded-by-tree-canopy=12
    target = target.to(torch.uint8)

rgb = image[0, :3].permute(1, 2, 0).byte()
image.shape, image.dtype, target.shape, target.dtype, rgb.shape, rgb.dtype

Visualize Input Image

Before inference, we inspect the RGB channels from the 4-band NAIP tensor (R, G, B, NIR). This confirms image orientation and content.

plt.figure(figsize=(4, 4))
plt.imshow(rgb)
plt.axis('off')
plt.title('NAIP RGB Tile')
plt.tight_layout()
plt.show()

Load Pre-trained Model

We load Unet_Weights.NAIP_RGBN_RESNET18_CHESAPEAKERSC and initialize unet with the matching number of classes.

Using the weights object ensures the model architecture and preprocessing transforms stay aligned with training configuration.

weights = Unet_Weights.NAIP_RGBN_RESNET18_CHESAPEAKERSC
model = unet(weights=weights).eval()
weights.meta

Run Inference

We apply the weights-provided transforms, run the model in inference mode, convert logits to class probabilities, and threshold road probabilities at 0.5 for a binary mask.

preprocess = weights.transforms
inputs = preprocess(image)
with torch.inference_mode():
    logits = model(inputs)
pred = logits.softmax(dim=1)[0, 1]
pred_mask = pred > 0.5
target_mask = target > 0
pred.shape, pred_mask.float().mean().item(), target_mask.float().mean().item()

Visualize Ground Truth vs Prediction

We compare:

• RGB input,

• ground-truth road mask,

• predicted road probability map,

• thresholded predicted mask.

This gives a quick qualitative check of segmentation behavior on a real sample tile.

fig, ax = plt.subplots(1, 4, figsize=(14, 4))
ax[0].imshow(rgb)
ax[0].set_title('RGB')
ax[1].imshow(target_mask, cmap='gray')
ax[1].set_title('Ground Truth Mask')
ax[2].imshow(pred, cmap='magma')
ax[2].set_title('Road Probability')
ax[3].imshow(pred_mask, cmap='gray')
ax[3].set_title('Predicted Mask > 0.5')
for a in ax:
    a.axis('off')
plt.tight_layout()
plt.show()