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

1. Enunciato del Problema

La segmentazione delle immagini mediche è fondamentale per la diagnosi clinica, ma incontra ostacoli significativi con i modelli fondazionali all'avanguardia attuali come il Segment Anything Model (SAM) e il suo successore SAM2.

• Divario di Dominio: SAM e SAM2 sono addestrati su immagini naturali e faticano a generalizzare sui dati medici a causa delle differenze nella texture, nel contrasto e nelle strutture anatomiche.
• Dipendenza dai Prompt: Questi modelli richiedono prompt forniti dall'utente (punti, riquadri o testo) per funzionare, il che è poco pratico per i flussi di lavoro clinici automatizzati.
• Scarsità di Dati: Il fine-tuning supervisionato richiede grandi dataset annotati da esperti, la cui creazione è costosa e dispendiosa in termini di tempo.
• Limitazioni delle Adattazioni Esistenti: I precedenti tentativi di adattare SAM/SAM2 ai dati medici spesso richiedono un fine-tuning computazionalmente oneroso o dipendono ancora dai prompt. Inoltre, le varianti senza prompt esistenti non sfruttano le capacità di memoria volumetrica di SAM2.

2. Metodologia: DiffuSAM

Gli autori propongono DiffuSAM, un framework che adatta SAM2 per la segmentazione medica senza prompt utilizzando un modello di diffusione latente condizionale. L'idea centrale è sintetizzare "embedding di memoria" che SAM2 possa utilizzare per generare maschere, sostituendo efficacemente la necessità di prompt utente.

Architettura Principale

1. Backbone Congelato: L'encoder di immagine e l'encoder di memoria di SAM2 rimangono congelati durante l'addestramento. Il modello apprende solo un prior di diffusione leggero.
2. Addestramento del Prior di Diffusione:
   • Input: Il modello prende l'embedding dell'immagine ($z_{img}$) dall'encoder di immagine congelato di SAM2 e l'etichetta di classe ($l$).
   • Obiettivo: Impara a generare il corrispondente embedding di memoria ($z_{mem}$) che SAM2 produrrebbe normalmente se gli venisse fornito un prompt.
   • Processo: Un modello di diffusione basato su UNet viene addestrato secondo una schedulazione del Denoising Diffusion Probabilistic Model (DDPM). Impara a rimuovere il rumore da un embedding di memoria rumoroso condizionato all'embedding dell'immagine e all'etichetta.
   • Funzione di Perdita: L'obiettivo di addestramento combina una perdita prior (ricostruzione dell'embedding di memoria) e una perdita di segmentazione (assicurando che l'embedding generato produca maschere accurate quando passato attraverso il decoder di maschere di SAM2).
3. Inferenza (Generazione Senza Prompt):
   • All'inferenza, il modello di diffusione inizia con puro rumore e l'embedding dell'immagine.
   • Rimuove iterativamente il rumore per sintetizzare un embedding di memoria ($\hat{z}_{mem}$).
   • Questo embedding sintetizzato viene iniettato nell'attenzione di memoria e nel decoder di maschere di SAM2 per produrre la maschera di segmentazione finale senza alcun input utente.

Coerenza Volumetrica (Condizionamento 3D)

Per affrontare la natura 3D dei volumi medici (TC/RM), DiffuSAM impone una coerenza spaziale tra le fette:

• Quando si genera l'embedding di memoria per una specifica fetta $i$, il modello si condiziona sulle previsioni di segmentazione delle fette adiacenti ($adj$).
• Le previsioni delle fette adiacenti vengono codificate tramite l'encoder di memoria e fuse con i dati della fetta corrente.
• Ciò garantisce la continuità anatomica e previene incoerenze tra le fette nel volume 3D.

Adattamento di Dominio Non Supervisionato Senza Sorgente (SF-UDA)

Il framework è progettato per SF-UDA, dove il modello è addestrato su un dominio sorgente (es. TC) e testato su un dominio target (es. RM) senza accesso alle immagini sorgente durante l'adattamento al target.

• Il metodo sfrutta la proprietà di affinità delle caratteristiche: gli embedding di immagine di SAM2 per la stessa classe tendono a raggrupparsi strettamente anche tra domini diversi.
• Il prior di diffusione addestrato sulla sorgente viene applicato direttamente alle fette del dominio target senza ulteriore fine-tuning, generando pseudo-embedding di memoria per il dominio target.

3. Contributi Chiave

1. Framework Innovativo: Introduzione di DiffuSAM, la prima adattamento basato su diffusione di SAM2 per la segmentazione di immagini mediche senza prompt.
2. Operazione nello Spazio Latente: Il modello opera nello spazio degli embedding latenti di SAM2, permettendogli di catturare distribuzioni di caratteristiche di alto livello senza modificare il pesante backbone.
3. Integrazione Volumetrica: Integrazione con il meccanismo di attenzione di memoria di SAM2 per garantire la continuità anatomica attraverso le fette 3D tramite condizionamento tra fette.
4. Efficienza: L'approccio richiede l'addestramento di un UNet leggero invece del fine-tuning dell'intero modello fondazionale, risultando in un basso utilizzo di memoria (<5GB GPU) e un'inferenza rapida (solo 2 passaggi di diffusione).

4. Risultati Sperimentali

Il metodo è stato valutato su due importanti dataset: BTCV (scansioni TC) e CHAOS (scansioni RM), focalizzandosi sulla segmentazione di milza, reni e fegato.

Impostazione Few-Shot (BTCV)

• Configurazione: Addestrato sul 10% dei dati (3 volumi), testato sul 90% (27 volumi).
• Prestazioni: DiffuSAM ha raggiunto un punteggio Dice medio del 87,24%.
• Confronto: Ha superato SAM/SAM2 nativi (che dipendono dai prompt) e modelli UNet/UNETR supervisionati addestrati sullo stesso piccolo dataset. È stato competitivo con FATE-SAM (che richiede fette ground-truth al momento del test) ma ha raggiunto questo risultato senza prompt manuali o supervisione al momento del test.
• Ablazione: L'aggiunta del condizionamento volumetrico tra fette ha migliorato il Dice medio dall'82,72% (DiffuSAM-no 3D) all'87,24%.

Impostazione SF-UDA (TC $\to$ RM)

• Configurazione: Addestrato su BTCV (TC), testato su CHAOS (RM) senza accesso alle immagini sorgente.
• Prestazioni: DiffuSAM ha raggiunto un Dice medio del 85,2%.
• Confronto: Ha superato significativamente altri baseline senza sorgente (DPL: 56,0%) e ha ottenuto prestazioni pari a DFG (85,2%), un forte metodo basato sull'affinità delle caratteristiche.
• Ablazione: Il condizionamento volumetrico ha migliorato le prestazioni dall'80,4% all'85,2%.

5. Significato e Conclusione

• Applicabilità Clinica: Eliminando la necessità di prompt utente, DiffuSAM colma il divario tra modelli fondazionali e flussi di lavoro clinici automatizzati.
• Efficienza dei Dati: Il metodo dimostra che è possibile ottenere segmentazioni di alta qualità con dati di addestramento minimi (few-shot) e senza accesso ai dati sorgente durante l'adattamento al target (SF-UDA).
• Efficienza Computazionale: A differenza del fine-tuning completo di grandi modelli, DiffuSAM addestra un prior di diffusione leggero, rendendolo accessibile per ambienti con risorse limitate.
• Lavori Futuri: Gli autori riconoscono che, sebbene MedSAM (addestrato end-to-end con prompt) performi meglio, DiffuSAM offre un vantaggio unico evitando prompt e fine-tuning completo. I lavori futuri esploreranno l'integrazione di ulteriori tecniche SF-UDA.

In sintesi, DiffuSAM sfrutta con successo il potere generativo dei modelli di diffusione per "allucinare" gli embedding di memoria necessari per SAM2, abilitando una segmentazione di immagini mediche robusta, senza prompt e adattabile al dominio.