Pixel-level Counterfactual Contrastive Learning for Medical Image Segmentation

1. Il Problema

La segmentazione delle immagini mediche dipende tradizionalmente da grandi dataset annotati manualmente, la cui produzione è costosa e lenta. Sebbene esistano etichette "silver-standard" (generate dall'AI), queste rischiano di introdurre bias. L'apprendimento auto-supervisionato (SSL) è diventato cruciale per il pre-addestramento utilizzando solo immagini non etichettate.

Tuttavia, le attuali tecniche di apprendimento contrastivo presentano due limiti principali nel contesto della segmentazione medica:

1. Livello di granularità: I metodi classici (es. SimCLR, MoCo) operano a livello di immagine, generando un singolo embedding globale, insufficiente per compiti di previsione densa come la segmentazione pixel-per-pixel.
2. Limiti dell'augmentazione: Le tecniche di augmentazione standard (rotazione, crop, sfocatura) non riescono a catturare gli spostamenti complessi dei dati (dataset shifts) tipici delle applicazioni reali, come variazioni nel tipo di scanner o nella presenza di patologie specifiche (es. versamento pleurico).

2. Metodologia

Gli autori propongono una pipeline che integra la generazione controfattuale con l'apprendimento contrastivo denso (pixel-level).

Generazione Controfattuale

Utilizzando un modello causale (Structural Causal Model - SCM) implementato tramite un Hierarchical Variational AutoEncoder (HVAE), il sistema genera variazioni realistiche e semanticamente significative delle immagini.

• Variabili causali: Tipo di scanner, sesso del paziente e presenza di versamento pleurico (PE).
• Interventi: Vengono generati tre tipi di controfattuali per ogni immagine base: cambio solo dello scanner, cambio solo della presenza di PE, o cambio di entrambi.
• Augmentazione: A questi controfattuali si aggiungono trasformazioni geometriche standard (rotazione, crop) ed effetti visivi.

Metodi di Apprendimento Contrastivo Denso

Vengono proposti quattro nuovi metodi che combinano le viste controfattuali con l'obiettivo contrastivo a livello di pixel:

1. DVD-CL (Dual-View Dense Contrastive Learning): Estende il framework VADeR. Una vista non controfattuale funge da "ancora", mentre le altre viste controfattuali sono "target". Le coppie positive sono pixel nella stessa posizione spaziale tra ancoraggio e target; tutte le altre combinazioni sono negative. La perdita (NT-Xent) è calcolata in entrambe le direzioni e mediata.
2. MVD-CL (Multi-View Dense Contrastive Learning): Estende il concetto a tutte le viste disponibili simultaneamente, ottimizzando le similarità tra tutte le coppie di pixel attraverso l'intero set di viste.
3. Varianti Supervise (S-DVD-CL e S-MVD-CL): Utilizzano le annotazioni silver-standard per definire le coppie positive e negative. I pixel della stessa classe (es. polmone sinistro) sono positivi, quelli di classi diverse (es. polmone destro vs sfondo) sono negativi. Questo approccio guida l'apprendimento verso una segmentazione più precisa.

Visualizzazione: CHRO-map

Viene introdotto un nuovo algoritmo di visualizzazione, la Color-coded High Resolution Overlay map (CHRO-map).

• Proietta gli embedding ad alta dimensionalità in uno spazio 2D usando UMAP.
• Mappa l'ellisse di volume minimo contenente i dati su un cerchio unitario.
• Assegna colori HSV basati sulle coordinate polari (angolo = tonalità, raggio = valore).
• Sovrappone il colore al pixel originale, permettendo una valutazione qualitativa della coerenza semantica e dell'invarianza controfattuale.

3. Contributi Chiave

• Quattro nuovi metodi: DVD-CL, MVD-CL e le loro varianti supervisionate (S-DVD-CL, S-MVD-CL), specifici per compiti di segmentazione densa.
• Integrazione controfattuale: Prima esplorazione dell'uso di aumentazioni controfattuali (basate su SCM) per l'apprendimento di rappresentazioni invarianti in contesti di segmentazione medica.
• CHRO-map: Uno strumento innovativo per visualizzare e interpretare gli embedding a livello di pixel, facilitando l'analisi della qualità delle rappresentazioni.
• Pre-addestramento completo: A differenza di lavori precedenti che pre-addestrano solo l'encoder, questo metodo pre-addestra l'intera struttura encoder-decoder di una U-Net.

4. Risultati

Gli esperimenti sono stati condotti sul dataset PadChest (circa 60k immagini di training) con annotazioni silver-standard CheXmask, e validati su un subset di 70 immagini manualmente annotate (20 sane, 50 con PE).

• Analisi dello Spazio Latente:
  • Le varianti non supervisionate (DVD-CL) mostrano una buona capacità di raggruppamento spaziale (i polmoni occupano regioni distinte) e invarianza rispetto allo scanner, anche se senza confini netti.
  • Le varianti supervisionate (S-DVD-CL, S-MVD-CL) apprendono cluster ben definiti corrispondenti alle classi anatomiche (polmone sinistro, destro, sfondo).
• Performance di Segmentazione (Fine-tuning):
  • Non supervisionato: DVD-CL supera i baselines esistenti (SimCLR, VADeR), raggiungendo un DSC (Dice Score) del 92.76%.
  • Supervisionato: Le varianti con etichette silver-standard migliorano ulteriormente i risultati. S-MVD-CL ottiene il miglior risultato assoluto con un DSC del 93.93% (e ~94% su dati difficili con PE), superando persino un modello pre-addestrato direttamente sulle etichette silver-standard (CheXMask, 93.65%).
  • I metodi proposti mostrano una maggiore robustezza alle variazioni patologiche (versamento pleurico) e di acquisizione rispetto ai metodi tradizionali.

5. Significato e Implicazioni

Questo lavoro dimostra che l'integrazione del ragionamento controfattuale nell'apprendimento contrastivo denso è un approccio promettente per sviluppare modelli di segmentazione medica più robusti e interpretabili.

• Robustezza: L'uso di controfattuali permette al modello di imparare caratteristiche invarianti rispetto a fattori di confusione (come il tipo di scanner o la presenza di malattie), migliorando la generalizzazione.
• Efficienza delle Etichette: Le varianti supervisionate dimostrano che l'uso intelligente di etichette "silver-standard" all'interno di un framework contrastivo può superare l'addestramento diretto su tali etichette, massimizzando il valore dei dati disponibili.
• Interpretabilità: La CHRO-map offre un nuovo modo per diagnosticare visivamente come il modello codifica le informazioni spaziali e semantiche, un aspetto cruciale per l'adozione clinica dell'AI.

In sintesi, il paper propone un framework che colma il divario tra la generazione di dati causali realistici e l'apprendimento di rappresentazioni dense, offrendo una soluzione scalabile e robusta per la segmentazione medica in scenari con dati limitati o annotazioni imperfette.