MedShift: Implicit Conditional Transport for X-Ray Domain Adaptation

1. Problema e Contexto

O uso de dados médicos sintéticos para treinar modelos robustos é uma solução escalável, mas enfrenta um grande desafio: a lacuna de domínio (domain gap) entre imagens sintéticas e reais. No contexto de radiografias de crânio, essa discrepância manifesta-se principalmente em três áreas:

• Comportamento de Atenuação: Simulações frequentemente não replicam a complexa atenuação de raios-X através de estruturas anatômicas heterogêneas (osso, cavidades de ar, tecidos moles).
• Características de Ruído: Sistemas reais introduzem ruído estruturado, espalhamento (scatter) e artefatos de compressão ausentes nos simuladores.
• Representação de Tecidos Moles e Contraste: Diferenças na dinâmica de contraste, especialmente nas fronteiras ósseas e em características anatômicas sobrepostas.

A maioria das abordagens atuais de tradução de imagem (Image-to-Image Translation) exige pares de dados alinhados (raramente disponíveis na medicina) ou requer o treinamento de um modelo separado para cada par de domínios, limitando a generalização e a escalabilidade.

2. Metodologia: MedShift

O artigo propõe o MedShift, um modelo generativo condicional unificado baseado em Flow Matching (FM) e Pontes de Schrödinger (Schrödinger Bridges).

• Espaço Latente Compartilhado: O modelo aprende um espaço latente agnóstico ao domínio. Durante a inferência, uma imagem de origem (ex: sintética) é codificada em uma representação latente $z_\tau$ através de integração temporal reversa (de $t=1$ até um tempo intermediário $\tau$).
• Transporte Condicional Implícito: Para gerar a imagem no domínio de destino (ex: real), o modelo realiza uma amostragem de tempo forward a partir de $z_\tau$ até $t=1$, condicionada ao rótulo do domínio alvo. Isso permite a tradução entre qualquer par de domínios vistos durante o treinamento sem necessidade de pares de dados.
• Arquitetura: O MedShift opera em um espaço latente aprendido por um VAE pré-treinado, reduzindo custos computacionais enquanto preserva detalhes semânticos. Utiliza um U-Net personalizado menor em comparação com modelos baseados em Difusão (DDPM) tradicionais.
• Flexibilidade na Inferência: O parâmetro $\tau$ (tempo intermediário) atua como um controle de trade-off:
  • $\tau$ alto: Preserva mais a estrutura anatômica da imagem original.
  • $\tau$ baixo: Aumenta a transformação estilística para o domínio alvo (mais realismo, mas risco de distorção estrutural).
• Guia Livre de Classificador (CFG): Utilizado para melhorar a aderência ao domínio alvo sem necessidade de treinamento adicional.

3. Contribuições Principais

1. Novo Modelo (MedShift): Uma abordagem baseada em Flow Matching e Pontes de Schrödinger para tradução de imagens não pareadas em alta resolução, capaz de generalizar entre múltiplos domínios em um único framework.
2. Novo Dataset (X-DigiSkull): Introdução de um dataset público contendo radiografias de crânio sintéticas e reais, alinhadas espacialmente, com múltiplos ângulos de aquisição e variações de dose de radiação (Baixa, Normal e Exposição).
3. Benchmarking Abrangente: Avaliação comparativa de vários modelos de ponta (GANs, Normalizing Flows, DDPMs) no contexto de adaptação de domínio médico, estabelecendo uma referência para trabalhos futuros.

4. Resultados e Desempenho

Os experimentos foram realizados na tarefa de converter raios-X sintéticos de alta dose para raios-X reais de dose normal.

• Qualidade vs. Estrutura: O MedShift alcançou o melhor equilíbrio entre realismo (medido por CFID e cobertura de distribuição) e preservação estrutural (medido por SSIM e LPIPS).
  • Modelos como CycleGAN-Turbo tiveram excelente realismo, mas introduziram distorções anatômicas significativas.
  • Hierarchy Flow preservou a estrutura, mas falhou em adaptar o estilo (domínio), resultando em pouca adaptação.
  • O MedShift, ajustando o parâmetro $\tau$, superou os concorrentes no ranking médio, oferecendo imagens que parecem reais sem destruir a anatomia subjacente.
• Eficiência Computacional:
  • O MedShift é 6 vezes menor em tamanho de modelo que os U-Nets do Stable Diffusion utilizados por concorrentes (SDEdit, Z-STAR).
  • Consome significativamente menos memória de vídeo (VRAM), tornando-o viável para implantação em hardware com restrições, embora exija múltiplos passos de inferência (diferente da inferência de um único passo de GANs).
• Qualidade Visual: As imagens traduzidas apresentaram transições suaves nas bordas (imitando o desfoque de sistemas reais), recuperação de detalhes de tecidos moles (como o couro cabeludo) e variações de intensidade internas mais naturais, superando a aparência "afiada" e artificial das simulações originais.

5. Significado e Conclusão

O MedShift representa um avanço significativo na adaptação de domínio para imagens médicas, oferecendo uma solução escalável e generalizável que não depende de dados pareados. Ao aprender um transporte implícito condicional, o modelo permite que dados sintéticos sejam utilizados de forma mais eficaz para treinar modelos clínicos, reduzindo a necessidade de grandes volumes de dados reais rotulados.

A introdução do dataset X-DigiSkull preenche uma lacuna importante na pesquisa, fornecendo um benchmark padronizado para avaliação de métodos de tradução entre raios-X sintéticos e reais. A capacidade de sintonizar o modelo para priorizar fidelidade perceptual ou consistência estrutural na inferência torna-o uma ferramenta prática e versátil para aplicações como navegação cirúrgica e treinamento de clínicos.