Context-Aware Semantic Segmentation via Stage-Wise Attention

1. Problema e Motivação

A segmentação semântica de imagens aéreas de ultra-alta resolução (UHR) é crucial para aplicações como mapeamento e monitoramento ambiental. No entanto, existem desafios significativos ao aplicar modelos baseados em Transformers a essas imagens:

• Complexidade Quadrática: A atenção global em Transformers cresce quadraticamente com o número de tokens, tornando o processamento de imagens UHR (que possuem milhões de pixels) proibitivo em termos de memória de GPU.
• Compromissos Atuais: Soluções comuns, como downsampling (redução de resolução) ou tiling (divisão em blocos), resultam na perda de contexto espacial de longo alcance ou na degradação de detalhes finos e bordas, prejudicando a precisão tanto em nível de objeto quanto de cena.
• Necessidade: É necessário um método que preserve detalhes locais finos (bordas, objetos pequenos) enquanto integra contexto global amplo, sem explodir o custo computacional.

2. Metodologia: CASWiT

Os autores propõem o CASWiT (Context-Aware Stage-Wise Transformer), uma arquitetura de Transformer hierárquico de dupla ramificação (dual-branch) projetada para segmentação UHR apenas com dados RGB.

Arquitetura Principal

• Dupla Ramificação:
  1. Ramo de Alta Resolução (HR): Processa crops (recortes) de alta resolução para preservar bordas e estruturas finas.
  2. Ramo de Baixa Resolução (LR): Processa patches mais amplos e reduzidos para codificar o contexto global da cena.
• Fusão por Atenção Cruzada em Etapas (Stage-Wise Cross-Attention):
  • Diferente de métodos anteriores que fundem características apenas no final (late fusion), o CASWiT injeta informações de contexto do ramo LR no ramo HR em todas as etapas do codificador (Stages 1 a 4).
  • Em cada etapa, as características HR atuam como Queries (consultas), enquanto as características LR fornecem Keys e Values (chaves e valores).
  • Isso permite uma injeção de contexto leve e eficiente desde os estágios iniciais da rede, mantendo a eficiência computacional.
• Decodificadores: O ramo HR utiliza cabeças de decodificação padrão como UPerNet ou SegFormer para gerar os logits finais. O ramo LR possui uma cabeça auxiliar usada apenas durante o treinamento para supervisão.

Estratégia de Pré-treinamento (SimMIM)

Para fortalecer a aprendizagem entre escalas, os autores adaptam uma estratégia de Modelagem de Imagem com Máscara (SimMIM):

• Mascaramento Duplo: Durante o pré-treinamento não supervisionado em ortofotos não rotuladas:
  • O ramo HR recebe mascaramento aleatório (75%).
  • O ramo LR recebe mascaramento centralizado (50%), cobrindo a área alinhada com o crop HR.
• Objetivo: Isso força a rede a reconstruir os pixels HR ausentes utilizando o contexto circundante do ramo LR, evitando a cópia trivial de pixels e incentivando a aprendizagem de correspondências cruzadas entre escalas.

3. Contribuições Chave

1. Arquitetura CASWiT: Introdução de um backbone dual-branch que funde detalhes HR com contexto LR através de atenção cruzada leve em múltiplas etapas do encoder, preservando detalhes finos.
2. Pré-treinamento SimMIM Adaptado: Desenvolvimento de uma estratégia de pré-treinamento específica para fluxo duplo que melhora a transferência e a robustez em tarefas UHR.
3. Protocolo de Avaliação RGB-Only UHR: Estabelecimento de um protocolo rigoroso no conjunto de dados FLAIR-HUB que utiliza a estrutura geoespacial para reconstruir grandes tiles contíguos, permitindo avaliação justa de métodos apenas RGB.
4. Generalização: Demonstração de que a arquitetura transfere eficazmente para outros domínios de alta resolução, especificamente segmentação médica.

4. Resultados Experimentais

Os experimentos foram realizados principalmente no conjunto de dados FLAIR-HUB e URUR, com validação em benchmarks médicos (ISIC e CRAG).

• FLAIR-HUB (Protocolo UHR RGB):

  • O CASWiT com decodificador SegFormer atingiu 66,37% de mIoU (média de Intersecção sobre União).
  • Este resultado supera as bases oficiais RGB (Swin-B + UPerNet) e até configurações multimodais anteriores (que usavam dados além de RGB).
  • Houve uma melhoria significativa na qualidade das bordas, com o mBIoU (Boundary IoU) subindo de 32,57 para 36,90.
  • O modelo é mais eficiente computacionalmente (298 GFLOPs) comparado a variantes com UPerNet.

• URUR (Benchmark Legado):

  • Alcançou 49,2% de mIoU, superando métodos específicos para UHR anteriores como WSDNet e Boosting Dual-Stream.

• Transferência para Medicina (ISIC e CRAG):

  • O modelo generalizou bem sem modificações arquiteturais, alcançando 86,5% de mIoU no ISIC e 90,7% no CRAG, superando bases de referência de dupla ramificação.

• Estudos de Ablação:

  • A fusão em todas as etapas (Stage-wise) superou a fusão apenas no início ou apenas no final.
  • A supervisão auxiliar no ramo LR (peso $\alpha=0,5$) foi crucial para regularizar a representação compartilhada.
  • O pré-treinamento SimMIM trouxe ganhos adicionais, especialmente na nitidez das bordas.

5. Significado e Conclusão

O trabalho demonstra que a injeção explícita e escalonada de contexto global em características de alta resolução é fundamental para a segmentação de imagens aéreas de ultra-alta resolução.

• Inovação: O CASWiT resolve o dilema entre resolução e contexto ao usar uma atenção cruzada eficiente em vez de fusão tardia ou downsampling agressivo.
• Impacto: O modelo estabelece um novo estado da arte (SOTA) para segmentação UHR apenas RGB, superando métodos que utilizam dados multimodais.
• Versatilidade: A eficácia do modelo em domínios médicos sugere que a abordagem de "contexto-aware" é uma solução geral para problemas de visão computacional que exigem alta fidelidade espacial e compreensão de cena ampla.

Em resumo, o CASWiT oferece uma arquitetura robusta e eficiente que equilibra a preservação de detalhes finos com a compreensão de contexto global, superando as limitações de memória dos Transformers tradicionais em cenários de ultra-alta resolução.