TerraX

TerraX extracts GIS features from high-resolution drone orthophotos. It is a multi-task segmentation model. One forward pass over a 512 × 512 RGB tile returns seven maps, one prediction per pixel:

#	Output	Type	What it finds
1	building_mask	binary	building footprints
2	roof_type_mask	5 classes	roof material of each building pixel: background, RCC (concrete), tiled, tin (metal sheet), other
3	road_mask	binary	road surface (area)
4	road_centerline_mask	binary, thin	road centre lines
5	waterbody_mask	binary	ponds, tanks and other water bodies (area)
6	waterbody_line_mask	binary, thin	narrow water channels such as canals, drains and streams
7	utility_line_mask	binary, thin	overhead utility lines such as power and telephone wires

TerraX was built for rural property and land-record mapping from drone surveys, such as India's SVAMITVA scheme. With the usual post-processing (threshold, clean-up, skeletonise, vectorise), the rasters become building polygons with a roof attribute, road polygons and centre lines, water polygons and lines, and utility lines. These can go straight into QGIS or ArcGIS.

---

Files

File	Size	Use it for
TerraX.onnx	271 MB	full-precision (FP32) model. Runs anywhere ONNX Runtime runs: CPU, CUDA, DirectML, CoreML
TerraX_light.onnx	136 MB	half-precision (FP16) weights, with float32 input and output. Half the size; agrees with TerraX.onnx on ≥ 99.9 % of pixels for every output
TerraX.pt	270 MB	PyTorch weights (model_state_dict), for fine-tuning or research

Both ONNX files take the same input and give the same outputs:

• Input image: float32, shape (batch, 3, 512, 512). The batch size is dynamic.
• Outputs: the seven names above, each float32, shape (batch, C, 512, 512). C = 5 for roof_type_mask and C = 1 for the rest. The values are logits: apply a sigmoid to the binary outputs and a softmax (or argmax) to roof_type_mask.

---

Quick start (ONNX, no PyTorch needed)

pip install onnxruntime numpy pillow huggingface_hub

import numpy as np, onnxruntime as ort
from PIL import Image
from huggingface_hub import hf_hub_download

path = hf_hub_download("aaronr43210/TerraX", "TerraX_light.onnx")   # or "TerraX.onnx"
sess = ort.InferenceSession(path, providers=["CPUExecutionProvider"])

MEAN = np.array([0.485, 0.456, 0.406], np.float32)
STD  = np.array([0.229, 0.224, 0.225], np.float32)

def preprocess(rgb):                       # rgb: (512, 512, 3) uint8 or float
    x = rgb.astype(np.float32)
    lo, hi = np.percentile(x, (2, 98))     # one stretch over the whole tile (all bands together)
    if hi - lo < 1e-6:
        hi = lo + 1.0
    x = np.clip((x - lo) / (hi - lo), 0, 1)
    x = (x - MEAN) / STD
    return x.transpose(2, 0, 1)[None].astype(np.float32)   # (1, 3, 512, 512)

tile = np.asarray(Image.open("tile_512.png").convert("RGB"))
names = [o.name for o in sess.get_outputs()]
out = dict(zip(names, sess.run(None, {"image": preprocess(tile)})))

sigmoid = lambda z: 1 / (1 + np.exp(-z))
building = sigmoid(out["building_mask"][0, 0]) > 0.5
roads    = sigmoid(out["road_mask"][0, 0]) > 0.5
roof     = out["roof_type_mask"][0].argmax(0)      # 0 bg, 1 RCC, 2 tiled, 3 tin, 4 other
roof[~building] = 0

To use a GPU, install onnxruntime-gpu and pass providers=["CUDAExecutionProvider", "CPUExecutionProvider"].

Preprocessing must match training

The model was trained on tiles prepared exactly like this:

1. Take the RGB bands only. Drop alpha and any NIR band.
2. Contrast-stretch the whole tile between its 2nd and 98th percentile, using all bands together, and clip to [0, 1].
3. Normalise each channel with the ImageNet mean (0.485, 0.456, 0.406) and std (0.229, 0.224, 0.225).

Skipping step 2 or 3 is the most common reason for empty or noisy predictions.

---

Running on a whole orthomosaic

Orthomosaics are far larger than 512 px, so predict on overlapping tiles and blend the results. The training pipeline used a 2D Gaussian window (σ = 128 px, ¼ of the tile) for this, which hides the seams at tile edges:

# pip install rasterio
import numpy as np, rasterio
from rasterio.windows import Window

T, OVERLAP = 512, 192
STRIDE = T - OVERLAP
g = np.exp(-0.5 * ((np.arange(T) - T / 2 + 0.5) / (T / 4)) ** 2)
KERNEL = np.outer(g, g).astype(np.float32)

with rasterio.open("orthophoto.tif") as src:
    H, W = src.height, src.width
    acc = {n: np.zeros(((5 if n == "roof_type_mask" else 1), H, W), np.float32) for n in names}
    wsum = np.zeros((H, W), np.float32)
    for y in range(0, H, STRIDE):
        for x in range(0, W, STRIDE):
            rgb = src.read([1, 2, 3], window=Window(x, y, T, T), boundless=True, fill_value=0)
            rgb = rgb.transpose(1, 2, 0)
            if not rgb.any():
                continue                                    # skip empty NoData tiles
            h, w = min(T, H - y), min(T, W - x)
            res = sess.run(None, {"image": preprocess(rgb)})
            for n, r in zip(names, res):
                p = (1 / (1 + np.exp(-r[0]))) if n != "roof_type_mask" else \
                    np.exp(r[0]) / np.exp(r[0]).sum(0, keepdims=True)
                acc[n][:, y:y + h, x:x + w] += p[:, :h, :w] * KERNEL[:h, :w]
            wsum[y:y + h, x:x + w] += KERNEL[:h, :w]
    prob = {n: a / np.maximum(wsum, 1e-8) for n, a in acc.items()}
    profile = src.profile                                  # keeps the CRS and geotransform for writing GeoTIFFs

Afterwards, threshold each probability map. For the thin outputs (centre lines, water lines, utility lines), skeletonise before vectorising. Then polygonise (for example with rasterio.features.shapes) and simplify. Use the source raster's transform so the vectors keep real-world coordinates.

---

PyTorch

TerraX.pt holds a model_state_dict for the TerraX architecture. To load it you need the model definition from the TerraX source code (models/model.py, segformer_encoder.py, decoder.py, heads.py) and transformers < 5. Version 5 renamed SegFormer's internal layers, so the keys no longer match.

import torch
from models.model import EnsembleDUKModel          # from the TerraX source code

model = EnsembleDUKModel(pretrained=False)
ckpt = torch.load("TerraX.pt", map_location="cpu")
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
with torch.no_grad():
    outputs = model(torch.randn(1, 3, 512, 512))   # dict with the seven outputs

Use the PyTorch weights if you want to fine-tune on your own region. For inference only, the ONNX files are simpler.

---

Architecture

RGB tile (3 × 512 × 512)
      │
SegFormer-B4 encoder (Mix Transformer, 60.8 M params)
      │   4 feature maps at 1/4, 1/8, 1/16, 1/32 resolution  (64, 128, 320, 512 channels)
      │
FPN decoder + CBAM attention at every pyramid level (3.5 M params)
      │   fused 256-channel feature map
      │
Seven task heads (3.1 M params), each output upsampled back to 512 × 512
  ├─ BuildingHead → building_mask + roof_type_mask (shared building features)
  ├─ BinaryHead   → road_mask, waterbody_mask
  └─ LineHead     → road_centerline_mask, waterbody_line_mask, utility_line_mask
                    (D-LinkNet-style dilated blocks for long, thin structures)

• Total: 67.5 M parameters.
• Encoder: SegFormer-B4, initialised from nvidia/segformer-b4-finetuned-cityscapes-1024-1024. Its overlapping patch embeddings and hierarchical self-attention capture context across a whole tile, which helps it follow long roads and lines.
• CBAM (Convolutional Block Attention Module) re-weights channels and spatial positions at each pyramid level, so small, thin targets are not lost among large background areas.
• Shared building features: the building head predicts the footprint and the roof material from the same features, so roof type is learned only where buildings are.

---

Training

Setting	Value
Data	Drone orthomosaics of rural villages from the SVAMITVA programme in India, from two regions (northern and central-eastern India); the full dataset has 9 orthomosaics at about 3.0–3.5 cm/px ground resolution
Labels	Vector annotations (shapefiles) for buildings with roof type, roads, road centre lines, water bodies, water lines and utility lines, rasterised to masks for each tile
Tiling	512 × 512 tiles with overlapping windows (about 35,000 tiles across all 9 orthomosaics); NoData pixels are masked out of the loss and the metrics
Validation split	by orthomosaic (whole maps held out) in the default configuration, so overlapping tiles do not leak between training and validation
Optimiser	AdamW, weight decay 1e-4 (none on norm and bias parameters)
Learning rate	one-cycle schedule with cosine annealing, peak 3e-4, momentum (β1) cycling 0.85–0.95
Precision	mixed precision (FP16/FP32), gradient clipping 0.5
Fine-tuning	whole network trained end to end; the same peak learning rate for every parameter group
Seed	42

Augmentation: random 90° rotations, flips and transposes; random resized crops; small affine and grid distortions; brightness, contrast, hue, saturation and RGB shifts; CLAHE; sharpening; Gaussian and median blur; Gaussian noise; coarse dropout; synthetic shadows.

Loss (per binary task): weighted sum of BCE (with online hard-example mining over the hardest 70 % of pixels), Dice, Tversky (α = 0.3, β = 0.7, which penalises missed pixels), focal (α = 0.25, γ = 2), Lovász hinge and a boundary-weighted BCE:

0.6·BCE + 0.9·Dice + 0.7·Tversky + 0.3·Focal + 1.0·Lovász + 0.5·Boundary

Roof type: cross-entropy over building pixels plus multi-class Dice. The task losses are weighted and then averaged over the tasks present in the batch.

This checkpoint

This release is the checkpoint saved at epoch 18 of a run planned for about 150 epochs. When it was saved it was the best epoch so far by validation mean IoU. It is an early checkpoint: area classes (buildings, water bodies) are already strong, while the thin-line classes are still under-trained.

---

Validation results (this checkpoint)

These are pixel-level metrics on held-out validation orthomosaics, with NoData excluded and threshold 0.5.

Task	IoU	Dice / F1	Precision	Recall
Building	88.0	93.6	92.5	94.8
Water body	83.9	91.2	88.6	94.0
Road (area)	68.0	80.9	72.9	91.0
Road centre line	65.6	79.2	70.9	89.8
Water line	38.8	55.9	65.5	48.8
Utility line	19.6	32.7	43.6	26.2
Mean (6 tasks)	60.6	72.3

Roof type (pixel accuracy on building pixels): 95.0 % overall. By class: RCC 97.8 %, tin 93.0 %, tiled 79.5 %, other 54.9 %.

The thin-line rows come from a pixel-overlap score. A line that is off by only a few pixels scores low even though its vector is nearly right, so for lines a buffer-based F1 (counted correct within a small distance) is the fairer measure.

---

What works well and what doesn't

Strong: building footprints, roof material on concrete and tin roofs, ponds and tanks, roads that are clearly visible.

Weaker, so check the output:

• Utility lines and narrow water channels. They are only a few pixels wide and often hidden by trees; expect gaps and missed segments.
• Roads under tree canopy come out broken. A road-graph clean-up step after prediction helps.
• Shadows and wet soil can be confused with water.
• "Other" and tiled roofs are confused more often than RCC or tin.
• Different imagery (another country, satellite data, or resolution far from 3–5 cm/px) has not been tested. Resample to about 3–3.5 cm/px, or fine-tune.

TerraX is a mapping aid, not a replacement for a surveyor. For cadastral or legal use, a person should review the extracted vectors.

---

Intended use

• Speeding up digitisation of rural drone surveys (buildings, roads, water, utilities).
• A first-pass vector layer for land-record, planning and infrastructure-inventory work.
• Research on multi-task remote-sensing segmentation.

Out of scope: surveillance or identifying people, and any decision about property rights made without human checking.

---

Citation

If you use TerraX, please cite:

@misc{terrax2026,
  title  = {TerraX: A Multi-Task Deep Learning Framework for Automated GIS Feature Extraction
            from High-Resolution Drone Orthophotos},
  author = {TerraX authors},
  year   = {2026},
  howpublished = {\url{https://huggingface.co/aaronr43210/TerraX}}
}