Context-Aware Semantic Segmentation via Stage-Wise Attention

1. Il Problema

La segmentazione semantica di immagini a ultra-alta risoluzione (UHR), tipiche del telerilevamento (es. mappatura aerea, monitoraggio ambientale), presenta una sfida fondamentale: è necessario preservare i dettagli locali fini (bordi, piccoli oggetti) mentre si sfrutta un contesto spaziale globale ampio.
I modelli basati su Transformer (come i ViT o Swin Transformer) hanno raggiunto lo stato dell'arte in molte vision tasks, ma la loro applicazione a input UHR è limitata dalla complessità quadratica dell'attenzione self-attention. Questo comporta un consumo di memoria GPU che cresce quadraticamente con il numero di token, costringendo spesso a compromessi dannosi:

• Downsampling: Riduce la risoluzione spaziale, perdendo dettagli critici.
• Tiling (finestre): Riduce il contesto globale, danneggiando la coerenza a livello di scena.
  Le soluzioni attuali spesso combinano elaborazione ad alta risoluzione e modellazione del contesto, ma molte si basano su fusioni tardive (mid/late fusion) che non permettono un'interazione efficiente tra le scale fin dalle prime fasi dell'encoding.

2. Metodologia: CASWiT

Gli autori propongono CASWiT (Context-Aware Stage-Wise Transformer), un'architettura a doppio ramo (dual-branch) basata su Swin Transformer, progettata specificamente per la segmentazione UHR RGB-only.

Architettura Principale

• Doppio Encoder:
  • Ramo ad Alta Risoluzione (HR): Elabora "crop" ad alta risoluzione (es. 512x512) per preservare i bordi e gli oggetti piccoli.
  • Ramo a Bassa Risoluzione (LR): Elabora patch più ampie ma a risoluzione inferiore (derivate da un campo visivo più grande) per codificare il contesto globale.
• Fusione Stage-Wise Cross-Attention:
  • A differenza delle fusioni tardive, CASWiT inietta il contesto LR nel flusso HR in ogni stadio dell'encoder (dallo stadio 1 al 4).
  • Utilizza blocchi di cross-attention leggeri: le feature HR agiscono come Query, mentre le feature LR agiscono come Key e Value.
  • Il risultato dell'attenzione viene fuso con le feature HR originali tramite una connessione residuale (e opzionalmente un gate appreso, sebbene la versione senza gate abbia mostrato prestazioni leggermente superiori).
• Decoder:
  • Utilizza un decoder UPerNet o SegFormer solo sul ramo HR per generare la mappa di segmentazione finale.
  • Il ramo LR ha un decoder ausiliario utilizzato solo durante l'addestramento per supervisione aggiuntiva.

Strategia di Pretraining (SimMIM)

Per migliorare l'apprendimento cross-scale, gli autori adattano una strategia di masked image modeling (stile SimMIM) su un vasto corpus di ortofoto non etichettate:

• Mascheramento HR: Applicazione casuale di maschere (75%) sul ramo HR.
• Mascheramento LR: Applicazione di una maschera centrata (50%) sul ramo LR, allineata spazialmente con il crop HR. Questo impedisce al modello di copiare trivialmente i pixel centrali dal LR all'HR, costringendolo a utilizzare il contesto circostante per ricostruire i pixel mancanti.
• L'obiettivo è ricostruire i pixel HR mascherati utilizzando sia le feature HR visibili che il contesto semantico LR.

3. Contributi Chiave

1. Architettura CASWiT: Introduzione di un backbone dual-branch gerarchico che inietta contesto LR nelle feature HR tramite cross-attention stage-wise, mantenendo l'efficienza computazionale.
2. Strategia di Pretraining: Sviluppo di un protocollo di pretraining self-supervised (SimMIM) specifico per flussi doppi, che rafforza l'apprendimento delle corrispondenze cross-scale.
3. Protocollo di Valutazione RGB-only UHR: Definizione di un protocollo rigoroso su FLAIR-HUB che ricostruisce grandi piastrelle contigue sfruttando la struttura geospaziale, permettendo una valutazione equa del contesto a lungo raggio senza dati multimediali aggiuntivi.
4. Generalizzazione: Dimostrazione che l'architettura non è limitata al telerilevamento, ma si trasferisce efficacemente alla segmentazione di immagini mediche UHR.

4. Risultati Sperimentali

I risultati sono stati ottenuti su dataset di grandi dimensioni come FLAIR-HUB, URUR, e benchmark medici (ISIC, CRAG).

• FLAIR-HUB (Protocollo UHR RGB-only):
  • CASWiT con decoder SegFormer raggiunge 66.37% mIoU, superando i baseline ufficiali RGB (Swin-B + UPerNet a 64.05%) e anche configurazioni multimodali precedenti (65.9% mIoU).
  • Migliora significativamente la qualità dei bordi: il mBIoU sale da 32.57 (baseline) a 36.90.
  • Riduce il costo computazionale rispetto a varianti con decoder UPerNet (298 GFLOPs vs 489 GFLOPs) aumentando la velocità di inferenza.
• URUR:
  • Raggiunge 49.2% mIoU, superando metodi specifici per UHR precedenti come WSDNet e Boosting Dual-Stream.
• Transfer Learning (Medicina):
  • Su ISIC (pelle) e CRAG (retina), CASWiT ottiene 86.5% e 90.7% mIoU rispettivamente, dimostrando che l'iniezione di contesto stage-wise è efficace anche in domini diversi dal telerilevamento.

5. Significato e Impatto

Il lavoro di CASWiT è significativo perché risolve il dilemma classico della segmentazione UHR: risoluzione vs. contesto.

• Efficienza: Dimostra che è possibile integrare il contesto globale senza dover processare l'intera immagine ad alta risoluzione con un singolo Transformer, evitando l'esplosione della memoria.
• Qualità dei Bordi: L'attenzione stage-wise permette di affinare i dettagli fini (bordi di edifici, strade strette) utilizzando il contesto semantico globale, riducendo il "bleeding" semantico tra classi adiacenti.
• Versatilità: La capacità di trasferire il modello a domini medici senza modifiche architetturali suggerisce che l'approccio è fondamentale per qualsiasi task di densa predizione su immagini ad altissima risoluzione.
• Open Source: Il codice e i modelli pre-addestrati sono stati rilasciati, facilitando la riproducibilità e l'adozione nella comunità di ricerca.

In sintesi, CASWiT stabilisce un nuovo stato dell'arte per la segmentazione semantica RGB-only su immagini aeree ad ultra-alta risoluzione, combinando un'architettura di fusione intelligente con strategie di pretraining avanzate.