Brain-SAM: A SAM-based Model Tailored for Brain MRI Lesion Segmentation

Resumo Técnico do Brain-SAM

Problema
A ressonância magnética (RM) é fundamental para a neuroimagem, mas a segmentação precisa de estruturas e lesões cerebrais permanece um desafio. Embora modelos de fundação como o Segment Anything Model (SAM) e seu sucessor SAM2 tenham demonstrado uma generalização notável em imagens naturais, sua aplicação à RM cerebral enfrenta obstáculos significativos. As adaptações médicas existentes (por exemplo, MedSAM, SAM-Med3D) são tipicamente treinadas em conjuntos de dados de órgãos mistos, contendo estruturas como coração, fígado e pulmões. Esses modelos têm dificuldade em se adaptar às características anatômicas únicas do cérebro, como a organização complexa dos tecidos, fronteiras indistintas entre a substância cinzenta e branca e morfologias de lesões altamente diversas. Além disso, a maioria das adaptações baseadas em SAM converte volumes 3D em fatias 2D, negligenciando o contexto espacial entre as fatias. Adicionalmente, a dependência de prompts interativos manuais limita a escalabilidade, pois os clínicos podem ter dificuldade em fornecer prompts precisos para patologias sutis, como lesões epilépticas. Por fim, o decodificador de máscara original do SAM, projetado para imagens naturais 2D com bordas claras, frequentemente apresenta desempenho subótimo em dados de RM cerebral, que são caracterizados por bordas borradas e heterogeneidade de tecidos.

Metodologia
Os autores propõem o Brain-SAM, um modelo adaptado especificamente para a segmentação de lesões em RM cerebral 3D. A arquitetura estende o framework SAM2 com três modificações primárias:

1. Adaptação 3D do Codificador de Imagem: O codificador Hiera 2D original (um Vision Transformer hierárquico) é adaptado para processar dados volumétricos 3D diretamente. Isso envolve a extensão de quatro módulos principais:

   • Patch Embedding: Substituição de convoluções 2D por convoluções 3D para particionar entradas volumétricas em patches de voxels.
   • Position Embedding: Extensão de representações 2D para 3D para capturar relações espaciais através de profundidade, altura e largura.
   • Token Sequence: Achatamento de volumes 3D ao longo de todas as três dimensões para criar sequências de tokens mais longas, adequadas para atenção volumétrica.
   • Q-Pooling: Substituição de MaxPool2d por MaxPool3d para permitir o subamostragem de características através das dimensões de profundidade e espacial.

2. Decodificador Inspirado em UNETR: Para abordar as limitações do decodificador original do SAM no tratamento de texturas locais finas e detalhes de borda, os autores introduzem um decodificador inspirado no UNETR. Este componente integra a capacidade de modelagem global dos Transformers com a representação local de grão fino das CNNs via fusão de características multiescala. Diferente dos decodificadores UNETR padrão que exigem transformações de resolução, este design aproveita as características hierárquicas inerentes do codificador Hiera para fusão direta e upsampling progressivo. A arquitetura suporta uma saída de cabeça dupla: uma para segmentação totalmente automática e outra para gerar características intermediárias para segmentação baseada em prompts.

3. Estratégia de Treinamento e Fine-Tuning:

   • Pré-treinamento: O modelo é inicializado com pesos do SAM2 e passa por um estágio de pré-treinamento no conjunto de dados BraTS2021 (1.000 casos, 4 modalidades) exclusivamente para segmentação automática, a fim de estabelecer representações médicas robustas.
   • Fine-Tuning Baseado em Adaptadores: Para adaptar-se a múltiplos conjuntos de dados de lesões cerebrais (tumores, AVC, epilepsia) enquanto evita o esquecimento catastrófico, módulos de adaptador leves (arquitetura bottleneck com convoluções 3D) são inseridos antes de cada bloco Hiera. Durante o fine-tuning, o backbone permanece congelado e apenas os parâmetros compartilhados do adaptador são atualizados.
   • Mecanismo de Prompt: Prompts de pontos são usados como o principal método de interação, estendidos para 3D. A estratégia de geração de máscara é modificada: em vez de usar chaves atualizadas na operação de produto escalar, o design funde queries processadas com representações de características intermediárias do decodificador inspirado em UNETR.
   • Refinamento: Um mecanismo de refinamento de pós-processamento utiliza cliques do usuário para ajustar localmente as probabilidades da máscara via núcleos Gaussianos 3D.

Principais Contribuições

1. Extensão 3D do SAM2: O trabalho estende com sucesso o framework SAM2 de imagens naturais 2D para análise de imagens médicas 3D, permitindo o processamento volumétrico direto sem aproximação fatia por fatia.
2. Arquitetura de Decodificador Híbrido: A introdução de um decodificador inspirado em UNETR permite que o modelo suporte tanto segmentação totalmente automática quanto interativa dentro de uma arquitetura unificada, capturando efetivamente tanto detalhes locais finos quanto dependências contextuais globais.
3. Curadoria Abrangente de Dados: Os autores curaram um conjunto de dados de grande escala cobrindo tumores cerebrais, AVC e epilepsia, permitindo que o modelo suporte uma ampla gama de aplicações clínicas.

Resultos
Experimentos extensivos foram conduzidos em sete conjuntos de dados públicos (BraTS-PED, BraTS-SSA, BraTS-MET, BraTS-MEN, ISLES2022, ATLAS v2.0, OPMRI) cobrindo seis modalidades de RM. O Brain-SAM foi comparado com modelos específicos para tarefas (nnU-Net, UNETR, Swin UNETR) e modelos de fundação (SAM-Med3D, VISTA3D).

• Desempenho Geral: O Brain-SAM alcançou escores Dice superiores e distâncias de Hausdorff 95% (HD95) menores em comparação com todos os baselines. Especificamente, comparado ao nnU-Net, alcançou melhorias de Dice de 22% em lesões epilépticas, 9% em metástases cerebrais e 6% em meningiomas.
• Segmentação de Pequenas Lesões: O modelo demonstrou vantagens significativas na segmentação de pequenas lesões (ex: epilepsia), alcançando escores Dice de 15% a 18% maiores do que os fortes modelos de baseline. Manteve um melhor equilíbrio entre recall e precisão para alvos pequenos (<5.000 voxels) em comparação com outros métodos.
• Estudos de Ablação:
  • Substituir o backbone 3D-ViT pelo 3D-Hiera melhorou os escores Dice em até 14% para pequenas lesões e reduziu significativamente o HD95.
  • A integração do decodificador inspirado em UNETR aumentou o escore Dice médio em 8% em comparação com a versão sem ele.
  • A estratégia de fine-tuning baseada em adaptadores provou ser altamente eficiente em termos de parâmetros (39,4M de parâmetros ajustáveis vs. 72,02M para fine-tuning total) enquanto alcança desempenho comparável ao fine-tuning total.

Significância e Alegações
O artigo afirma que o Brain-SAM oferece um modelo pré-treinado robusto para tarefas de análise de RM cerebral downstream. Ao abordar os desafios anatômicos e estruturais específicos do tecido cerebral — como bordas indistintas e tamanhos de lesões pequenas — ele fornece segmentação mais confiável e precisa do que os modelos de fundação generalistas existentes. Os autores postulam que essa capacidade pode apoiar a tomada de decisão clínica mais bem informada, particularmente no diagnóstico precoce e planejamento de tratamento para condições como epilepsia e tumores cerebrais. A capacidade do modelo de funcionar tanto em modo automático quanto guiado por prompts visa reduzir a carga de trabalho de anotação manual, mantendo alta precisão. Os autores observam que o estudo atual é limitado a dados de modalidade única e sugerem que trabalhos futuros possam explorar a integração de múltiplas modalidades para aumentar ainda mais a caracterização das lesões.