Pixel-level Counterfactual Contrastive Learning for Medical Image Segmentation

1. Problema e Motivação

A segmentação de imagens médicas depende fortemente de grandes conjuntos de dados anotados, que são caros e demorados para produzir. Embora existam alternativas como:

• Rótulos "Silver-standard" (gerados por IA): Mais fáceis de obter, mas tendem a introduzir viés.
• Aprendizado Auto-supervisionado: Utiliza apenas imagens, sendo crucial para pré-treinamento.

No entanto, os métodos atuais de Aprendizado Contrastivo (que aprendem representações trazendo amostras similares próximas e afastando as dissimilares) têm limitações significativas:

1. Nível de Imagem vs. Pixel: Métodos clássicos (como SimCLR) operam no nível da imagem inteira, gerando um único embedding. Isso é inadequado para tarefas de previsão densa, como segmentação, que exigem informações espacialmente localizadas.
2. Limitações de Aumento de Dados: As transformações tradicionais (rotação, corte, desfoque) não capturam as mudanças complexas e desafiadoras de conjuntos de dados do mundo real (ex: variações de equipamentos de imagem ou presença de patologias).

O objetivo deste trabalho é superar essas limitações combinando Geração Contrafactual com Aprendizado Contrastivo Densa (Pixel-level) para melhorar a robustez e a generalização na segmentação de imagens médicas.

2. Metodologia

Os autores propõem um pipeline que integra a geração de imagens contrafactuais com objetivos contrastivos densos.

A. Geração Contrafactual

Utilizam um Hierarchical Variational AutoEncoder (HVAE) baseado em Modelos Causais Estruturais (SCM).

• Variáveis Causais: Tipo de scanner, sexo e presença de Derrame Pleural (PE).
• Processo: O modelo gera versões contrafactuais da imagem original intervindo em variáveis específicas (ex: mudar apenas o tipo de scanner, ou apenas a presença de PE, ou ambos), mantendo os outros fatores fixos. Isso simula cenários de "e se?".
• Aumento Adicional: As imagens geradas passam por aumentos tradicionais (rotação, corte, jitter de cor, etc.).

B. Métodos de Aprendizado Contrastivo Densa

São propostos quatro métodos que utilizam as visões geradas (imagem base + contrafactuais) para treinar um modelo U-Net (pré-treinando todo o encoder-decoder):

1. DVD-CL (Dual-View Dense Contrastive Learning):
   • Usa a imagem base como "âncora" e cada imagem contrafactual como "alvo".
   • Calcula a perda contrastiva (NT-Xent) entre pixels na mesma localização espacial entre a âncora e o alvo.
   • A perda é média entre todas as direções e todas as visões alvo.
2. MVD-CL (Multi-View Dense Contrastive Learning):
   • Considera todas as visões simultaneamente.
   • Calcula a perda sobre todos os pares de pixels entre todas as visões disponíveis, permitindo uma otimização conjunta.
3. Variantes Supervisionadas (S-DVD-CL e S-MVD-CL):
   • Utilizam anotações "silver-standard" (CheXmask) para definir pares positivos e negativos.
   • Pares Positivos: Pixels da mesma classe anatômica (ex: ambos pulmões esquerdos).
   • Pares Negativos: Pixels de classes diferentes (ex: pulmão esquerdo vs. direito) ou fundo.
   • Apenas pixels não-fundo são usados como âncoras para garantir supervisão significativa.

C. Visualização: CHRO-map

Foi introduzido um novo algoritmo de visualização chamado Color-coded High Resolution Overlay map (CHRO-map):

• Projeta embeddings de alta dimensão em 2D usando UMAP.
• Transforma a projeção em um círculo unitário e atribui cores (HSV) baseadas nas coordenadas polares (ângulo = matiz, raio = valor).
• Sobrepõe a cor ao pixel original, permitindo visualizar a similaridade semântica e a invariância contrafactual diretamente na imagem.

3. Contribuições Principais

1. Quatro novos métodos de aprendizado contrastivo denso: DVD-CL, MVD-CL e suas variantes supervisionadas (S-DVD-CL, S-MVD-CL).
2. Integração de Geração Contrafactual: Uso de SCMs para criar aumentos de dados semanticamente significativos e realistas, superando as limitações dos aumentos tradicionais.
3. CHRO-map: Uma ferramenta de visualização de alta resolução para interpretar embeddings de pixels e avaliar a qualidade da representação aprendida.
4. Exploração pioneira: Primeira aplicação de aprendizado contrastivo pixel-level com geração contrafactual para segmentação médica.

4. Resultados Experimentais

Os experimentos foram realizados no conjunto de dados PadChest (Chest X-Ray) com o objetivo de segmentar pulmões na presença de Derrame Pleural (PE).

• Análise do Espaço Latente (Pré-treinamento):

  • DVD-CL (Não supervisionado): Mostrou-se superior a outros métodos não supervisionados (SimCLR, VADeR), aprendendo representações que distinguem os pulmões e capturam a relação espacial com a coluna vertebral.
  • MVD-CL (Não supervisionado): Desempenhou-se pior, falhando em capturar estruturas semânticas claras.
  • Variantes Supervisionadas: Alcançaram alta pureza de agrupamento (K-means) e separação clara entre classes (pulmão esquerdo, direito e fundo) nos mapas CHRO.

• Avaliação de Fine-tuning (Segmentação):

  • Desempenho Geral: Todos os métodos propostos superaram suas respectivas linhas de base (baselines).
  • Não Supervisionado: O DVD-CL obteve o melhor desempenho entre os não supervisionados, com um DSC (Dice Score) médio de 92.76%, superando o VADeR (92.26%) e o SimCLR (90.90%).
  • Supervisionado: As variantes supervisionadas superaram o treinamento direto nos dados "silver-standard".
    • O S-MVD-CL alcançou o melhor resultado global, com um DSC de ~93.93% (arredondado para ~94% no resumo do artigo), superando o modelo pré-treinado diretamente no CheXmask (93.65%).
    • Houve melhoria significativa na robustez para pacientes com PE (DSCPE), onde o S-MVD-CL atingiu 92.43%.

5. Significado e Conclusão

O trabalho demonstra que a integração de raciocínio contrafactual no aprendizado contrastivo denso melhora significativamente a robustez e a interpretabilidade de modelos de segmentação médica.

• Robustez: O uso de contrafactuais (mudança de scanner e patologia) torna o modelo invariante a fatores de confusão, melhorando a generalização para dados desafiadores (como PE).
• Eficiência de Dados: As variantes supervisionadas mostram que, mesmo com anotações imperfeitas ("silver-standard"), o uso de aprendizado contrastivo com contrafactuais gera representações superiores ao treinamento supervisionado direto.
• Interpretabilidade: A ferramenta CHRO-map oferece uma nova maneira de validar qualitativamente se o modelo está aprendendo características anatômicas corretas e invariantes.

Limitações e Futuro: O método assume que o grafo causal é conhecido e que os contrafactuais são identificáveis a partir dos dados observados, o que pode não ser sempre verdade na prática. Trabalhos futuros visam explorar métodos de geração contrafactual alternativos e escalar o framework para modalidades 3D e cenários semi-supervisionados.