CASWiT

CASWiT: Context-Aware Stage Wise Transformer for Ultra-High Resolution Semantic Segmentation

Context-Aware Semantic Segmentation via Stage-Wise Attention

License: MIT SOTA: FLAIR HUB @ RGB SOTA: URUR SOTA: ISIC SOTA: CRAG

Official implementation of CASWiT, a dual-branch architecture for ultra-high resolution semantic segmentation that leverages stage-wise cross-attention fusion between high-resolution and low-resolution branches.

📋 Table of Contents

• Overview
• Architecture
• Installation
• Docker
• Dataset Preparation
• Usage
• Configuration
• Results
• Citation
• License

🎯 Overview

CASWiT addresses semantic segmentation on ultra-high resolution (UHR) imagery with a dual-resolution design:

• HR Branch: processes high-resolution crops (e.g., 512×512) for fine details
• LR Branch: processes low-resolution context (typically downsampled) for global information
• Stage-wise Cross-Attention Fusion: HR features attend to LR context at each encoder stage

CASWiT generalizes beyond aerial imagery without architectural changes, and we provide configs/scripts for several UHR datasets (FLAIR-HUB, URUR, SWISSIMAGE inference, and additional medical benchmarks).

🏗️ Architecture

CASWiT architecture

Key components:

• Dual Swin Transformer Backbones
• Cross-Attention Fusion Blocks at each encoder stage
• Auxiliary LR Supervision (optional, weighted by model.lr_supervision_weight)

📦 Installation

Requirements

• Python 3.12+
• PyTorch 2+
• CUDA 12+ (for GPU training)

Setup

git clone https://huggingface.co/heig-vd-geo/CASWiT
cd CASWiT
pip install -r requirements.txt

🐳 Docker

A Dockerfile is provided for a reproducible environment.

Build

docker build -t caswit:latest .

Run (GPU)

docker run --gpus all -it --rm \
  -v $(pwd):/workspace \
  caswit:latest

If your datasets/checkpoints are outside the repo, mount them too, e.g.:

docker run --gpus all -it --rm \
  -v $(pwd):/workspace \
  -v /path/to/data:/data \
  -v /path/to/checkpoints:/checkpoints \
  caswit:latest

📊 Dataset Preparation

FLAIR-HUB

1. Download the FlairHub dataset
2. Merge GeoTIFF tiles into mosaics:

python dataset/prepareFlairHub.py

URUR

Expected structure:

URUR/
├── train/
│   ├── image/
│   └── label/
├── val/
│   ├── image/
│   └── label/
└── test/
    ├── image/
    └── label/

A re-hosted copy is available here: https://huggingface.co/datasets/heig-vd-geo/URUR

SWISSIMAGE

Download images using the provided CSV:

python dataset/download_swissimage.py list_all_swiss_image_sept2025.csv

SwissImage + SwissTLM3D Adaptation Pipeline

The repository also includes a full adaptation pipeline for Switzerland using:

• SwissImage as imagery source
• SwissTLM3D as coarse / partial supervision
• FLAIR-HUB as the semantic anchor for mixed fine-tuning

1. Download SwissImage

python dataset/download_swissimage.py list_all_swiss_image_sept2025.csv

2. Download SwissTLM3D

The official SwissTLM3D GeoPackage can be downloaded from swisstopo:

mkdir -p /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d
curl -L \
  https://data.geo.admin.ch/ch.swisstopo.swisstlm3d/swisstlm3d_2026-02-24/swisstlm3d_2026-02-24_2056_5728.gpkg.zip \
  -o /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d/swisstlm3d_2026-02-24_2056_5728.gpkg.zip

unzip -o \
  /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d/swisstlm3d_2026-02-24_2056_5728.gpkg.zip \
  -d /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d

3. Rasterize SwissTLM3D masks on SwissImage

We recommend the conservative reliable5 rasterization variant used in our Swiss adaptation runs. It projects only a sparse high-confidence subset of classes and sets all unsupported pixels to 255:

• 0: building
• 5: bare soil
• 6: water (Stehende Gewaesser only)
• 11: vineyard
• 12: forest (Wald, Wald offen, Gehoelzflaeche)

The historical reliable5 script name is kept for compatibility, even though the current mapping is now a conservative 5-class Swiss adaptation subset.

python dataset/rasterize_swisstlm3d_reliable5.py \
  --gpkg /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d/SWISSTLM3D_2026_LV95_LN02.gpkg \
  --image_dir /mnt/CalcShare/datasets/Suisse_full/data/image \
  --output_dir /mnt/CalcShare/datasets/Suisse_full/data/mask_tlm3d_reliable5 \
  --workers 24

4. Create 2048 -> 1024 training patches

SwissImage has a coarser GSD than FLAIR-HUB in our setup. We therefore extract 2048x2048 windows and resize them to 1024x1024 before training.

For the sparse reliable5 supervision, we keep only patches whose central HR-supervised region contains at least 20% labeled pixels. Concretely, the ratio is computed on the central 1024x1024 region of the original 2048x2048 patch before resize, matching the central HR crop later supervised during training.

python dataset/prepare_swisstlm3d_patches_reliable5.py \
  --image_dir /mnt/CalcShare/datasets/Suisse_full/data/image \
  --mask_dir /mnt/CalcShare/datasets/Suisse_full/data/mask_tlm3d_reliable5 \
  --output_dir /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d_reliable5_forest_center20_patches_2048to1024 \
  --tile_size 2048 \
  --output_size 1024 \
  --min_labeled_ratio 0.20 \
  --workers 48

5. Create a random split by source tile

The current Swiss adaptation setup uses a deterministic random split at the source-tile level, so all patches extracted from the same original SwissImage tile stay in the same split.

python -m dataset.split_swisstlm3d_geographic_reliable5 \
  --image_dir /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d_reliable5_forest_center20_patches_2048to1024/img \
  --mask_dir /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d_reliable5_forest_center20_patches_2048to1024/msk \
  --output_dir /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d_reliable5_forest_center20_patches_2048to1024_split \
  --seed 42 \
  --overwrite

🚀 Usage

All commands below work either directly with the scripts in train/ or via the unified main.py.

Training

Single GPU:

python train/train.py configs/config_FlairHub.yaml

Multi-GPU (DDP):

torchrun --nproc_per_node=4 train/train.py configs/config_FlairHub.yaml

Evaluation

python train/eval.py configs/config_FlairHub.yaml weights/checkpoint.pth test

Inference (single image)

python train/inference.py configs/config_FlairHub.yaml weights/checkpoint.pth image.tif output.png

Inference on a VRT (DDP)

This runs tiled inference on a VRT using multiple GPUs:

torchrun --nproc_per_node=5 --master_port=29501 -m train.inference_vrt_ddp \
  --config configs/config_SWISSIMAGE_inf.yaml \
  --checkpoint weights/CASWiT-Base-SSL_FLAIRHUB_UN.pth \
  --vrt file.vrt \
  --out_dir output/ \
  --tile 1024 \
  --stride 512 \
  --lr_side 2048

Using main.py

# Train
python main.py train --config configs/config_FlairHub.yaml

# Eval (default split = test in the script if not specified)
python main.py eval --config configs/config_FlairHub.yaml --checkpoint weights/checkpoint.pth

# Inference (single image)
python main.py inference --config configs/config_FlairHub.yaml --checkpoint weights/checkpoint.pth --image image.tif --output pred.png

⚙️ Configuration

Configs are YAML files in configs/.

Model selection (decoder head)

You can select the decoder head directly in the config:

model:
  # options: upernet | segformer | mask2former | fusion_last_stage_add | ssl
  head: upernet

Notes:

• mask2former may require training.amp: false depending on your environment/precision settings.
• ssl is the SimMIM-style pretraining model and is not a drop-in replacement for segmentation inference/eval scripts.

Example config structure

paths:
  data_path: "/path/to/dataset"
  dataset_name: "FLAIRHUB"
  train_img_subdir: "train/img"
  train_msk_subdir: "train/msk"
  val_img_subdir: "valid/img"
  val_msk_subdir: "valid/msk"
  test_img_subdir: "test/img"
  test_msk_subdir: "test/msk"
  save_dir: "weights"
  pretrained_path: ""

model:
  model_name: "openmmlab/upernet-swin-base"
  num_classes: 15
  cross_attention_heads: 1
  ignore_index: 255
  fusion_mlp_ratio: 4.0
  fusion_drop_path: 0.1
  lr_supervision_weight: 0.5
  # options: upernet | segformer | mask2former | fusion_last_stage_add | ssl
  head: upernet

training:
  batch_size: 4
  num_workers: 8
  num_epochs: 20
  learning_rate: 0.00006
  amp: true
  seed: 42
  eta_min: 0.000001

wandb:
  use_wandb: true
  project: "CASWiT"
  entity: "your-entity"
  run_name: "caswit_experiment"

augmentations:
  enable: false
  p_hflip: 0.5
  p_vflip: 0.5
  p_rot90: 0.5
  color_jitter:
    brightness: 0.2
    contrast: 0.2
    saturation: 0.2
    hue: 0.05
  blur:
    p: 0.1
    kernel: 3

🧪 Reproducing experiments (all provided configs)

Below are ready-to-run commands for each config. Replace --checkpoint with your own file when evaluating/inferencing.

FLAIR-HUB

# Train (no extra aug)
python main.py train --config configs/config_FlairHub.yaml

# Train (with augmentations)
python main.py train --config configs/config_FlairHub_aug.yaml

# Eval
python main.py eval --config configs/config_FlairHub.yaml --checkpoint weights/checkpoint.pth

URUR

python main.py train --config configs/config_URUR.yaml
python main.py train --config configs/config_URUR_aug.yaml

python main.py eval --config configs/config_URUR.yaml --checkpoint weights/checkpoint.pth

ISIC (medical)

python main.py train --config configs/config_ISIC_aug.yaml
python main.py eval  --config configs/config_ISIC_aug.yaml --checkpoint weights/checkpoint.pth

CRAG (medical)

python main.py train --config configs/config_CRAG_aug.yaml
python main.py eval  --config configs/config_CRAG_aug.yaml --checkpoint weights/checkpoint.pth

Swiss Fine-Tuning

Mixed FLAIR-HUB + SwissTLM3D fine-tuning

This is the recommended strategy for adaptation:

• keep FLAIR-HUB as the strong semantic supervision source
• inject SwissTLM3D patches as domain adaptation supervision
• validate separately on FLAIR-HUB and SwissTLM3D
• select the best checkpoint using the configured mixed validation score

The current reference config is:

• configs/config_mixed_domain_ft_segformer_reliable5.yaml

It expects the reliable5 random split, e.g.:

paths:
  swiss_path: /mnt/CalcShare/datasets/Suisse_full/data/swisstlm3d_reliable5_forest_center20_patches_2048to1024_split

Then launch:

torchrun --nproc_per_node=4 --master_port=29501 \
  train/train_mixed_domain_ft.py \
  configs/config_mixed_domain_ft_segformer_reliable5.yaml

The mixed script logs:

• val_flair_mIoU, val_flair_mF1
• val_swiss_mIoU, val_swiss_mF1
• model_selection_score

The default mixed strategy uses:

• swiss_ratio: 0.20
• best_on: mean
• a 15-class head kept for exact...