Enhancing MedSAM with a Lightweight Box Predictor for Medical Image Segmentation

Resumo Técnico: Aprimorando o MedSAM com um Preditor de Caixa Leve para Segmentação de Imagens Médicas

Definição do Problema
A segmentação semântica em imagens médicas é dificultada pela escassez de dados, alta variabilidade entre modalidades (CT, MRI, Ultrassom) e a dificuldade de obter conjuntos de dados anotados em larga escala. Embora modelos de fundação como o Segment Anything Model (SAM) ofereçam capacidades de segmentação de uso geral, sua aplicação direta a imagens médicas é limitada por lacunas de domínio. Especificamente, a adaptação do MedSAM, embora melhorada, permanece altamente sensível à qualidade dos prompts do usuário. O prompt de ponto padrão, embora natural para a interação do usuário, frequentemente fornece contexto espacial insuficiente para uma segmentação confiável, particularmente quando as estruturas alvo são irregulares, pequenas ou possuem baixo contraste. Essa ambiguidade leva a uma degradação significativa do desempenho, especialmente quando o modelo não é totalmente ajustado (fine-tuned) ou quando os cliques do usuário são imprecisos.

Metodologia
Os autores propõem um framework de segmentação aprimorado que integra um módulo de Preditor de Caixa (Box Predictor) leve à arquitetura do MedSAM. A inovação central é a conversão de um único clique do usuário (prompt de ponto) em uma caixa delimitadora aproximada, que serve como um prior espacial robusto para guiar o processo de segmentação.

O método opera através de um pipeline de treinamento de dois estágios:

1. Estágio 1 (Treinamento do Preditor de Caixa): Um Perceptron de Camadas Múltiplas (MLP) leve é treinado independentemente para prever parâmetros de caixa delimitadora ($W, H, \Delta x, \Delta y$) a partir de um patch localizado de $5 \times 5$ de embeddings de imagem extraídos ao redor do ponto de entrada. O ponto de entrada é sujeito a perturbações espaciais aleatórias durante o treinamento (escalonadas para 30% da área média da máscara) para simular a imprecisão realista do usuário. O módulo é otimizado usando uma função de perda composta combinando perda $L1$ e perda de Interseção sobre União Generalizada (GIoU).
2. Estágio 2 (Integração e Ajuste Fino): O Preditor de Caixa pré-treinado é integrado ao pipeline do MedSAM. O Codificador de Imagem e o Preditor de Caixa são mantidos congelados para preservar as representações aprendidas. O Codificador de Prompt e o Decodificador de Máscara podem ser opcionalmente ajustados (fine-tuned). A caixa delimitadora prevista é alimentada no Codificador de Prompt juntamente com o prompt de ponto original (ou independentemente) para gerar a máscara de segmentação final.

O Preditor de Caixa adiciona apenas 1,6 milhão de parâmetros e introduz um overhead de inferência desprezível (aproximadamente 61–63 ms por imagem em uma RTX 4090).

Principais Contribuições

• Preditor de Caixa Leve: Um módulo que transforma um prompt de ponto único em uma caixa delimitadora informativa, fornecendo orientação espacial que estabiliza o MedSAM, particularmente quando os prompts de ponto sozinhos são ambíguos ou mal posicionados.
• Pipeline de Treinamento de Dois Estágios: Uma estratégia onde o Preditor de Caixa é treinado separadamente e depois integrado como orientação espacial fixa. Isso permite que o sistema alcance uma segmentação robusta mesmo quando o modelo MedSAM principal permanece congelado, tornando-o adequado para ambientes com recursos limitados.
• Avaliação Sistemática: Avaliação abrangente em quatro conjuntos de dados heterogêneos (FLARE22, BRISC, BUSI, LungSegDB) cobrindo as modalidades CT, MRI e Ultrassom. O estudo avalia seis configurações, analisando a interação entre tipos de prompt (ponto, caixa, híbrido) e estados do modelo (congelado vs. ajustado).
• Análise de Robustez de Prompt: Investigação de como perturbações de entrada e ajustes no tamanho da caixa afetam a precisão, demonstrando que o método permanece eficaz mesmo com cliques descentralizados e requer ajuste manual mínimo para diferentes modalidades.

Resultados Experimentais
O método proposto foi avaliado contra o baseline MedSAM e outros variantes do SAM (SAM2, SAM2UNet) em quatro conjuntos de dados. As principais descobertas incluem:

• Ganhos de Desempenho: O método melhorou consistentemente a precisão da segmentação. No conjunto de dados BUSI, o score Dice aumentou de 0,8771 (MedSAM + ponto) para 0,8904 (MedSAM + Caixa). No FLARE22, o score Dice melhorou de 0,9169 para 0,9314.
• Recuperação de Modelo Congelado: Em um cenário "congelado" onde os pesos do MedSAM não foram atualizados, confiar apenas em prompts de ponto resultou em falha quase total (scores Dice tão baixos quanto 1,64% no LungSegDB). No entanto, a introdução da caixa delimitadora prevista restaurou significativamente o desempenho (por exemplo, o Dice subiu para 72,58% no LungSegDB), demonstrando que o prior geométrico é crítico para a operação do modelo.
• Adaptação de Modalidade: O método mostrou comportamentos distintos entre as modalidades. Conjuntos de dados de Ultrassom e MRI (BUSI, BRISC) beneficiaram-se de um leve aumento na caixa delimitadora (10–20%) para compensar fronteiras ambíguas, enquanto conjuntos de dados de CT (FLARE22, LungSegDB) performaram de forma ideal com as caixas previstas iniciais devido ao maior contraste.
• Eficiência: A abordagem alcançou esses ganhos com custo computacional mínimo, mantendo tempos de inferência comparáveis ao baseline.

Significância e Alegações
O artigo afirma que este trabalho não redefine o paradigma fundamental de segmentação, mas oferece um aprimoramento prático ao MedSAM. Sua principal significância reside na redução da dependência de prompts manuais precisos, melhorando assim a robustez da segmentação interativa em fluxos de trabalho clínicos reais. Ao alavancar um módulo auxiliar leve, o método preenche a lacuna entre as expectativas geométricas do modelo de fundação e as limitações do input esparso do usuário. Os autores enfatizam que esta abordagem é particularmente valiosa em cenários com dados rotulados ou recursos computacionais limitados, pois permite uma segmentação de alto desempenho sem a necessidade de um ajuste fino completo do modelo. O estudo conclui que, embora o método tenha limitações em relação a regiões múltiplas e desconectadas, ele aborda efetivamente a ambiguidade dos prompts de ponto através de diversos domínios de imagem médica.