DiffuSAM: Diffusion-Based Prompt-Free SAM2 for Few-Shot and Source-Free Medical Image Segmentation

1. Formulação do Problema

A segmentação de imagens médicas é crítica para o diagnóstico clínico, mas enfrenta obstáculos significativos com os modelos fundamentais mais avançados atuais, como o Segment Anything Model (SAM) e seu sucessor SAM2.

• Lacuna de Domínio: O SAM e o SAM2 são treinados em imagens naturais e têm dificuldade em generalizar para dados médicos devido a diferenças em textura, contraste e estruturas anatômicas.
• Dependência de Prompts: Esses modelos exigem prompts fornecidos pelo usuário (pontos, caixas ou texto) para funcionar, o que é impraticável para fluxos de trabalho clínicos automatizados.
• Escassez de Dados: O ajuste fino supervisionado requer grandes conjuntos de dados anotados por especialistas, que são caros e demorados para criar.
• Limitações das Adaptações Existentes: Tentativas anteriores de adaptar o SAM/SAM2 a dados médicos frequentemente exigem um ajuste fino computacionalmente pesado ou ainda dependem de prompts. Além disso, as variantes sem prompts existentes não aproveitam as capacidades de memória volumétrica do SAM2.

2. Metodologia: DiffuSAM

Os autores propõem o DiffuSAM, um framework que adapta o SAM2 para segmentação médica sem prompts utilizando um modelo de difusão latente condicional. A ideia central é sintetizar "embeddings de memória" que o SAM2 pode usar para gerar máscaras, substituindo efetivamente a necessidade de prompts do usuário.

Arquitetura Central

1. Backbone Congelado: O codificador de imagem e o codificador de memória do SAM2 permanecem congelados durante o treinamento. O modelo aprende apenas um prior de difusão leve.
2. Treinamento do Prior de Difusão:
   • Entrada: O modelo recebe o embedding da imagem ($z_{img}$) do codificador de imagem congelado do SAM2 e o rótulo da classe ($l$).
   • Alvo: Aprende a gerar o embedding de memória correspondente ($z_{mem}$) que o SAM2 normalmente produziria se recebesse um prompt.
   • Processo: Um modelo de difusão baseado em UNet é treinado sob um agendamento de Modelo Probabilístico de Difusão com Remoção de Ruído (DDPM). Aprende a remover ruído de um embedding de memória com ruído, condicionado ao embedding da imagem e ao rótulo.
   • Função de Perda: O objetivo de treinamento combina uma perda de prior (reconstruindo o embedding de memória) e uma perda de segmentação (garantindo que o embedding gerado produza máscaras precisas quando passada pelo decodificador de máscara do SAM2).
3. Inferência (Geração sem Prompts):
   • Na inferência, o modelo de difusão começa com ruído puro e o embedding da imagem.
   • Remove ruído iterativamente para sintetizar um embedding de memória ($\hat{z}_{mem}$).
   • Este embedding sintetizado é injetado na atenção de memória e no decodificador de máscara do SAM2 para produzir a máscara de segmentação final sem qualquer entrada do usuário.

Consistência Volumétrica (Condicionamento 3D)

Para abordar a natureza 3D de volumes médicos (TC/RM), o DiffuSAM impõe consistência espacial entre fatias:

• Ao gerar o embedding de memória para uma fatia específica $i$, o modelo condiciona-se às previsões de segmentação de fatias adjacentes ($adj$).
• As previsões das fatias adjacentes são codificadas via codificador de memória e fundidas com os dados da fatia atual.
• Isso garante continuidade anatômica e previne inconsistências entre as fatias no volume 3D.

Adaptação de Domínio Não Supervisionada Livre de Fonte (SF-UDA)

O framework é projetado para SF-UDA, onde o modelo é treinado em um domínio de origem (por exemplo, TC) e testado em um domínio de destino (por exemplo, RM) sem acesso às imagens de origem durante a adaptação ao destino.

• O método aproveita a propriedade de afinidade de características: os embeddings de imagem do SAM2 para a mesma classe tendem a agrupar-se estreitamente, mesmo entre domínios diferentes.
• O prior de difusão treinado na origem é aplicado diretamente às fatias do domínio de destino sem ajuste fino adicional, gerando pseudo-embeddings de memória para o domínio de destino.

3. Contribuições Principais

1. Framework Inovador: Introdução do DiffuSAM, a primeira adaptação baseada em difusão do SAM2 para segmentação médica sem prompts.
2. Operação no Espaço Latente: O modelo opera no espaço de embeddings latentes do SAM2, permitindo capturar distribuições de características de alto nível sem modificar o backbone pesado.
3. Integração Volumétrica: Integração com o mecanismo de atenção de memória do SAM2 para garantir continuidade anatômica entre fatias 3D via condicionamento entre fatias.
4. Eficiência: A abordagem requer o treinamento de uma UNet leve em vez de ajustar todo o modelo fundamental, resultando em baixo uso de memória (<5GB de GPU) e inferência rápida (apenas 2 passos de difusão).

4. Resultados Experimentais

O método foi avaliado em dois conjuntos de dados principais: BTCV (tomografias computadorizadas) e CHAOS (ressonâncias magnéticas), focando na segmentação do baço, rins e fígado.

Configuração Few-Shot (BTCV)

• Configuração: Treinado em 10% dos dados (3 volumes), testado em 90% (27 volumes).
• Desempenho: O DiffuSAM alcançou uma pontuação média Dice de 87,24%.
• Comparação: Superou o SAM/SAM2 nativo (que dependem de prompts) e modelos supervisionados UNet/UNETR treinados no mesmo pequeno conjunto de dados. Foi competitivo com o FATE-SAM (que requer fatias de verdade no momento do teste), mas alcançou isso sem prompts manuais ou supervisão no momento do teste.
• Ablação: Adicionar condicionamento volumétrico entre fatias melhorou o Dice médio de 82,72% (DiffuSAM-no 3D) para 87,24%.

Configuração SF-UDA (TC $\to$ RM)

• Configuração: Treinado no BTCV (TC), testado no CHAOS (RM) sem acesso às imagens de origem.
• Desempenho: O DiffuSAM alcançou um Dice médio de 85,2%.
• Comparação: Superou significativamente outras bases livres de origem (DPL: 56,0%) e performou na paridade com o DFG (85,2%), um forte método baseado em afinidade de características.
• Ablação: O condicionamento volumétrico melhorou o desempenho de 80,4% para 85,2%.

5. Significado e Conclusão

• Aplicabilidade Clínica: Ao eliminar a necessidade de prompts do usuário, o DiffuSAM preenche a lacuna entre modelos fundamentais e fluxos de trabalho clínicos automatizados.
• Eficiência de Dados: O método demonstra que segmentação de alta qualidade pode ser alcançada com dados de treinamento mínimos (few-shot) e sem acesso a dados de origem durante a adaptação ao destino (SF-UDA).
• Eficiência Computacional: Diferente do ajuste fino completo de grandes modelos, o DiffuSAM treina um prior de difusão leve, tornando-o acessível para ambientes com recursos limitados.
• Trabalho Futuro: Os autores reconhecem que, embora o MedSAM (treinado de ponta a ponta com prompts) tenha melhor desempenho, o DiffuSAM oferece uma vantagem única ao evitar prompts e o ajuste fino completo. O trabalho futuro explorará a integração de técnicas adicionais de SF-UDA.

Em resumo, o DiffuSAM aproveita com sucesso o poder generativo dos modelos de difusão para "alucinar" os embeddings de memória necessários para o SAM2, permitindo segmentação de imagens médicas robusta, sem prompts e adaptativa a domínios.