Enhancing MedSAM with a Lightweight Box Predictor for Medical Image Segmentation

Sintesi Tecnica: Potenziamento di MedSAM con un Predittore di Bounding Box Leggero per la Segmentazione di Immagini Mediche

Problematica
La segmentazione semantica nelle immagini mediche è ostacolata dalla scarsità di dati, dall'elevata variabilità tra le modalità (CT, MRI, Ultrasuoni) e dalla difficoltà di ottenere dataset annotati su larga scala. Sebbene i modelli di fondazione come il Segment Anything Model (SAM) offrano capacità di segmentazione generaliste, la loro applicazione diretta alle immagini mediche è limitata dai gap di dominio. Nello specifico, l'adattamento di MedSAM, pur essendo migliorato, rimane altamente sensibile alla qualità dei prompt dell'utente. Il tipico prompt a punto (point prompt), sebbene naturale per l'interazione dell'utente, spesso fornisce un contesto spaziale insufficiente per una segmentazione affidabile, particolarmente quando le strutture target sono irregolari, piccole o presentano un basso contrasto. Questa ambiguità porta a un significativo degrado delle prestazioni, specialmente quando il modello non è completamente sottoposto a fine-tuning o quando i clic dell'utente sono imprecisi.

Metodologia
Gli autori propongono un framework di segmentazione potenziato che integra un modulo leggero di Box Predictor nell'architettura MedSAM. L'innovazione principale è la conversione di un singolo clic dell'utente (point prompt) in un bounding box approssimativo, che funge da robusto prior spaziale per guidare il processo di segmentazione.

Il metodo opera tramite una pipeline di addestramento a due stadi:

1. Stadio 1 (Addestramento del Box Predictor): Un leggero Multi-Layer Perceptron (MLP) viene addestrato indipendentemente per predire i parametri del bounding box ($W, H, \Delta x, \Delta y$) da una patch locale di embedding d'immagine $5 \times 5$ estratta attorno al punto di input. Il punto di input è soggetto a perturbazione spaziale casuale durante l'addestramento (scalata al 30% dell'area media della maschera) per simulare l'imprecisione realistica dell'utente. Il modulo è ottimizzato utilizzando una funzione di perdita composita che combina la perdita $L1$ e la perdita Generalized Intersection over Union (GIoU).
2. Stadio 2 (Integrazione e Fine-tuning): Il Box Predictor pre-addestrato viene integrato nella pipeline MedSAM. L'Image Encoder e il Box Predictor vengono mantenuti congelati (frozen) per preservare le rappresentazioni apprese. Il Prompt Encoder e il Mask Decoder possono essere opzionalmente sottoposti a fine-tuning. Il bounding box predetto viene inserito nel Prompt Encoder insieme al punto originale (o indipendentemente) per generare la maschera di segmentazione finale.

Il Box Predictor aggiunge solo 1,6 milioni di parametri e introduce un overhead di inferenza trascurabile (circa 61–63 ms per immagine su una RTX 4090).

Contributi Chiave

• Box Predictor Leggero: Un modulo che trasforma un singolo prompt a punto in un bounding box informativo, fornendo una guida spaziale che stabilizza MedSAM, specialmente quando i prompt a punto sono ambigui o posizionati male.
• Pipeline di Addestramento a Due Stadi: Una strategia in cui il Box Predictor viene addestrato separatamente e poi integrato come guida spaziale fissa. Ciò consente al sistema di raggiungere una segmentazione robusta anche quando il modello MedSAM principale rimane congelato, rendendolo adatto ad ambienti con risorse limitate.
• Valutazione Sistematica: Valutazione completa su quattro dataset eterogenei (FLARE22, BRISC, BUSI, LungSegDB) che coprono le modalità CT, MRI e Ultrasuoni. Lo studio valuta sei configurazioni, analizzando l'interazione tra tipi di prompt (punto, box, ibrido) e stati del modello (congelato vs. fine-tuned).
• Analisi della Robustezza del Prompt: Investigazione su come le perturbazioni dell'input e le regolazioni della dimensione del box influenzino l'accuratezza, dimostrando che il metodo rimane efficace anche con clic fuori centro e richiede una minima regolazione manuale per diverse modalità.

Risultati Sperimentali
Il metodo proposto è stato valutato rispetto ai baseline MedSAM e altre varianti di SAM (SAM2, SAM2UNet) su quattro dataset. Le scoperte chiave includono:

• Guadagni di Prestazioni: Il metodo ha costantemente migliorato l'accuratezza della segmentazione. Sul dataset BUSI, il punteggio Dice è aumentato da 0,8771 (MedSAM + punto) a 0,8904 (MedSAM + Box). Su FLARE22, il punteggio Dice è migliorato da 0,9169 a 0,9314.
• Recupero del Modello Congelato: In una configurazione "frozen" in cui i pesi di MedSAM non sono stati aggiornati, fare affidamento esclusivamente sui prompt a punto ha portato a un quasi totale fallimento (punteggi Dice così bassi come l'1,64% su LungSegDB). Tuttavia, l'introduzione del bounding box predetto ha ripristinato significativamente le prestazioni (ad esempio, il Dice è salito al 72,58% su LungSegDB), dimostrando che il prior geometrico è critico per l'operatività del modello.
• Adattamento alla Modalità: Il metodo ha mostrato comportamenti distinti tra le modalità. I dataset di Ultrasuoni e MRI (BUSI, BRISC) hanno beneficiato di un leggero ingrandimento del bounding box (10–20%) per tenere conto dei confini ambigui, mentre i dataset CT (FLARE22, LungSegDB) hanno performato in modo ottimale con i box iniziali predetti grazie al contrasto più elevato.
• Efficienza: L'approccio ha ottenuto questi guadagni con un costo computazionale minimo, mantenendo tempi di inferenza comparabili al baseline.

Significatività e Rivendicazioni
Il documento afferma che questo lavoro non ridefinisce il paradigma fondamentale della segmentazione, ma offre un potenziamento pratico a MedSAM. La sua principale significatività risiede nella riduzione della dipendenza da prompt manuali precisi, migliorando così la robustezza della segmentazione interattiva nei flussi di lavoro clinici reali. Sfruttando un modulo ausiliario leggero, il metodo colma il divario tra le aspettative geometriche del modello di fondazione e i limiti del limitato input dell'utente. Gli autori sottolineano che questo approccio è particolarmente prezioso in contesti con dati etichettati o risorse computazionali limitate, poiché consente una segmentazione ad alte prestazioni senza la necessità di un pieno fine-tuning del modello. Lo studio conclude che, sebbene il metodo presenti dei limiti riguardo a regioni multiple e disconnesse, affronta efficacemente l'ambiguità dei prompt a punto attraverso diversi domini di imaging medico.