Domain and Task-Focused Example Selection for Data-Efficient Contrastive Medical Image Segmentation

1. Il Problema

La segmentazione delle immagini mediche è fondamentale per la diagnosi e il monitoraggio, ma l'approccio supervisionato completo richiede grandi quantità di dati annotati a livello di pixel. Questa annotazione è costosa, richiede tempo ed è soggetta a errori, limitando l'adozione di modelli efficaci in ambito clinico.
Le tecniche di apprendimento auto-supervisionato (SSL), in particolare l'apprendimento contrastivo, offrono una soluzione pre-addestrando modelli su dati non etichettati. Tuttavia, le strategie contrastive esistenti per le immagini mediche presentano limiti significativi:

• Dipendenza da aumentazioni dei dati: Spesso si basano su trasformazioni geometriche o di colore che non catturano le relazioni semantiche complesse presenti nei dati medici.
• Mancanza di contesto: Non sfruttano efficacemente le relazioni intrinseche tra diverse immagini (es. relazioni inter-scan o intra-scan) o la conoscenza specifica del dominio (es. presenza di un organo target).
• Generalizzazione: I modelli spesso faticano a generalizzare tra diversi domini o protocolli di acquisizione.

2. Metodologia: PolyCL

Gli autori propongono PolyCL, un nuovo framework di apprendimento contrastivo auto-supervisionato progettato specificamente per la segmentazione medica. Il framework segue un processo in due fasi:

A. Pre-addestramento Auto-Supervisionato

Invece di creare coppie positive tramite aumentazioni casuali, PolyCL utilizza strategie innovative di selezione degli esempi basate sul dominio e sul compito. Per ogni "slice" (fetta) di riferimento (ancora), vengono selezionati esempi positivi e negativi secondo tre varianti:

1. PolyCL-O (Organ-based): Richiede la conoscenza binaria della presenza dell'organo target. Le coppie positive condividono lo stesso stato (entrambe contengono l'organo o entrambe no), mentre le negative hanno stati opposti. Questo insegna al modello a riconoscere la struttura target.
2. PolyCL-S (Scan-based): Non richiede annotazioni aggiuntive. Gli esempi positivi provengono dalla stessa scansione (intra-scan), mentre quelli negativi provengono da scansioni diverse (inter-scan). Questo cattura la coerenza intra-scan e l'invarianza inter-scan.
3. PolyCL-M (Mixed): Combina le due strategie precedenti. Un esempio positivo deve provenire dalla stessa scansione e condividere lo stato di presenza dell'organo. Questo approccio ibrido mira a massimizzare sia la coerenza contestuale che la discriminazione semantica.

L'obiettivo di perdita è la classica contrastive loss (InfoNCE), che massimizza la similarità tra l'ancora e l'esempio positivo e minimizza quella con l'esempio negativo nello spazio latente.

B. Fine-tuning Supervisionato

Dopo il pre-addestramento, la testa di proiezione viene rimossa. Un decoder (architettura U-Net style) viene aggiunto all'encoder pre-addestrato e l'intero modello viene fine-tunato su un piccolo set di dati etichettati utilizzando una funzione di perdita basata sul coefficiente Dice.

C. Integrazione con Segment Anything Model (SAM)

Per migliorare ulteriormente i risultati con dati limitati, il framework integra il modello SAM in due modi innovativi:

1. Raffinamento delle Maschere (Post-processing): Le maschere grezze prodotte da PolyCL vengono convertite in bounding box. Queste box vengono usate come prompt per SAM, che rifinisce i contorni anatomici sfruttando le sue capacità di segmentazione zero-shot.
2. Propagazione delle Maschere (SAM 2): Sfruttando le capacità di segmentazione video di SAM 2, è possibile propagare una singola maschera annotata 2D attraverso l'intero volume 3D CT. Questo permette di generare segmentazioni volumetriche complete partendo da un'unica slice annotata.

3. Contributi Chiave

• Nuova Strategia di Pre-addestramento: Introduzione di PolyCL, che utilizza la selezione di esempi basata su organi e scansioni invece delle tradizionali aumentazioni di dati, permettendo al modello di apprendere rappresentazioni più rilevanti per il compito.
• Strategie di Selezione Ibride: Sviluppo di tre varianti (O, S, M) che bilanciano la necessità di annotazioni minime con la necessità di contesto semantico.
• Integrazione SAM Avanzata: Applicazione di SAM non solo per il raffinamento post-hoc, ma anche come meccanismo di propagazione 3D (tramite SAM 2) per ridurre drasticamente il carico di annotazione.
• Validazione Estensiva: Sperimentazione su tre dataset pubblici CT (LiTS, TotalSegmentator, MSD) e KiTS, dimostrando superiorità sia in scenari in-domain che cross-domain.

4. Risultati Sperimentali

Le valutazioni quantitative e qualitative mostrano che PolyCL supera i baseline supervisionati completi e altri metodi SSL (come SimCLR, SupCon, e metodi specifici per l'anatomia) in scenari con dati limitati:

• Performance In-Domain: Su LiTS (fegato) e KiTS (reni), le varianti di PolyCL (specialmente PolyCL-M e PolyCL-O) raggiungono coefficienti Dice competitivi o superiori rispetto ai modelli supervisionati completi, con miglioramenti significativi nella distanza di Hausdorff (migliore accuratezza dei bordi).
• Generalizzazione Cross-Domain: In test su dataset esterni (TotalSegmentator, MSD, BTCV) senza ri-addestramento specifico, PolyCL dimostra una capacità di generalizzazione superiore, mantenendo performance robuste dove i modelli supervisionati falliscono o degradano.
• Efficienza con Dati Limitati: Con solo il 5-10% di dati etichettati, PolyCL combinato con SAM (PolyCL+SAM) ottiene risultati superiori rispetto ai metodi di raffinamento esistenti (es. SAMRefiner).
• Propagazione 3D: L'uso di SAM 2 permette di ottenere segmentazioni volumetriche di alta qualità partendo da una singola slice annotata, riducendo il lavoro di annotazione umana in modo esponenziale.

5. Significato e Impatto

Il lavoro di PolyCL è significativo per diversi motivi:

• Riduzione del Carico di Annotazione: Dimostra che è possibile ottenere segmentazioni di alta qualità con una frazione minima di dati etichettati, rendendo la segmentazione medica più accessibile e scalabile.
• Adattabilità al Dominio: Le strategie di selezione degli esempi basate sul contesto medico (organi e scansioni) colmano il divario tra l'apprendimento contrastivo generico e le esigenze specifiche delle immagini mediche.
• Pipeline Clinica Efficiente: L'integrazione con SAM 2 offre una soluzione pratica per la segmentazione volumetrica rapida, cruciale per la pianificazione chirurgica e il monitoraggio, riducendo il tempo necessario ai radiologi.
• Flessibilità Architetturale: Il framework funziona efficacemente con diverse architetture di backbone (U-Net, ResU-Net, TransUNet), dimostrando la sua versatilità.

In sintesi, PolyCL rappresenta un avanzamento significativo verso l'efficienza dei dati nella segmentazione medica, combinando strategie di pre-addestramento intelligenti con le più recenti capacità di modelli foundation come SAM.