Context-Aware Semantic Segmentation via Stage-Wise Attention

TL;DR

The paper tackles RGB-only semantic segmentation of ultra-high-resolution aerial imagery, addressing the memory bottlenecks of transformers by introducing CASWiT, a dual-branch Swin-based architecture that injects global context into high-resolution features via stage-wise cross-attention. It couples a supervised objective with an auxiliary low-resolution loss and leverages a SimMIM-style self-supervised pretraining on unlabeled orthophotos to enhance cross-scale learning. Empirical results on FLAIR-HUB and URUR demonstrate state-of-the-art RGB-only performance and strong boundary quality, with SSL and augmentation yielding notable gains ($mIoU$ up to $65.83\%$ on FLAIR-HUB and $49.1\%$ on URUR). The work offers a practical RGB-only UHR segmentation solution with robust cross-scale representations and broad applicability in remote sensing tasks, along with public code and pretrained weights for reproducibility.

Abstract

Semantic ultra high resolution image (UHR) segmentation is essential in remote sensing applications such as aerial mapping and environmental monitoring. Transformer-based models struggle in this setting because memory grows quadratically with token count, constraining either the contextual scope or the spatial resolution. We introduce CASWiT (Context-Aware Stage-Wise Transformer), a dual-branch, Swin-based architecture that injects global cues into fine-grained UHR features. A context encoder processes a downsampled neighborhood to capture long-range dependencies, while a high resolution encoder extracts detailed features from UHR patches. A cross-scale fusion module, combining cross-attention and gated feature injection, enriches high-resolution tokens with context. Beyond architecture, we propose a SimMIM-style pretraining. We mask 75% of the high-resolution image tokens and the low-resolution center region that spatially corresponds to the UHR patch, then train the shared dual-encoder with small decoder to reconstruct the UHR initial image. Extensive experiments on the large-scale IGN FLAIR-HUB aerial dataset demonstrate the effectiveness of CASWiT. Our method achieves 65.83% mIoU, outperforming RGB baselines by 1.78 points. On URUR, CASWiT achieves 49.1% mIoU, surpassing the current SoTA by +0.9% under the official evaluation protocol. All codes are provided on: https://huggingface.co/collections/heig-vd-geo/caswit.

Figures (13)

• Figure 1: Proposed architecture (CASWiT). A dual-branch (green) encoder for ultra–high-resolution imagery: the high-resolution (HR) branch processes the targeted tile to predict, while the low-resolution (LR) branch (pink) ingests a downsampled global context. At every Swin stage (1$\rightarrow$4), HR features supply queries ($Q$) and LR features provide keys/values ($K,V$) to a multi-scale cross-attention module with gating ($\otimes$) and residual connections ($\oplus$), enabling controlled context exchange (right panel). Two decoder heads (HR/LR) are jointly optimized via a weighted cross-entropy loss ($\mathcal{L}$), injecting global LR cues while preserving HR detail. Intermediate resolutions ($H/4 \dots H/32$) follow Swin patch-merging; a MLP further refines the fused output.
• Figure 2: Self-supervised inference results on the CASWiT architecture. Each image (left to right) shows: original high-resolution image, high-resolution image with random masking, low-resolution image with central masking, and the reconstruction of the high-resolution image after SimMIM-style pretraining.
• Figure 3: HR/LR construction on FLAIR-HUB. Red: original HR patch (512$\times$512). Green: georeferenced 3$\times$3 neighborhood assembled into a 1024$\times$1024 context, then downsampled $\times$2 to form the LR input (512$\times$512).
• Figure 4: URUR: illustrative annotation mismatch. Example where the provided mask (overlaid) locally diverges from the RGB content; such cases are occasional but can affect evaluation metrics. See §. \ref{['sec:datasets']} and supplementary for more examples.
• Figure 5: Qualitative comparison on IGN FLAIR-HUB. From left to right: LR image (note the missing band at the top), HR image crop, ground-truth overlay, RGB baseline (Swin-Base) overlay, and CASWiT overlay. CASWiT better recovers small vegetation patches and yields crisper boundaries, closely matching the GT while reducing false positives on bare soil and road areas (bottom-left). Despite the LR artifact (black band), CASWiT remains stable.