Improving Prostate Gland Segmentation Using Transformer based Architectures

Titolo: Miglioramento della Segmentazione della Ghiandola Prostatica tramite Architetture Basate su Transformer

1. Il Problema

La segmentazione automatica della ghiandola prostatica nelle immagini MRI pesate in T2 (T2-weighted) è fondamentale per la diagnosi precoce e la pianificazione terapeutica del cancro alla prostata. Tuttavia, l'automazione di questo processo incontra due sfide principali:

• Variabilità inter-lettore: Le annotazioni manuali effettuate da diversi radiologi (esperti) presentano differenze significative (variabilità di stile di delineazione), il che rende difficile addestrare modelli generalizzabili.
• Squilibrio delle classi e dimensione ridotta: La prostata occupa meno del 3% del volume dell'immagine, creando uno sbilanciamento estremo tra foreground (prostata) e background.
• Limiti delle CNN tradizionali: Le Reti Neurali Convoluzionali (CNN), come la 3D U-Net, hanno un campo ricettivo locale e faticano a catturare le dipendenze a lungo raggio necessarie per segmentare strutture complesse e piccole in contesti eterogenei. Inoltre, i modelli basati su CNN tendono a generalizzare male quando testati su dati annotati da lettori diversi da quelli usati per l'addestramento.

2. Metodologia

Gli autori hanno condotto uno studio sistematico confrontando architetture basate su Transformer con un modello CNN di riferimento (3D U-Net) su un dataset di 546 volumi MRI T2 annotati indipendentemente da due esperti (Reader #1 e Reader #2).

Architetture Confrontate:

• 3D U-Net (Baseline): Utilizza blocchi convolutivi-pooling con connessioni skip.
• UNETR: Sostituisce l'encoder CNN con un Vision Transformer (ViT) puro. Divide il volume in patch non sovrapposte ($16^3$) e applica l'attenzione self-attention globale. Tuttavia, soffre di una crescita quadratica della memoria.
• SwinUNETR: Utilizza un encoder basato su Swin Transformer con meccanismo di finestre spostate (shifted-window). Questo approccio calcola l'attenzione all'interno di finestre locali ($7^3$) che si spostano tra i livelli, permettendo una crescita lineare della memoria e catturando il contesto globale in modo efficiente.

Strategie di Addestramento e Valutazione:
Lo studio ha implementato tre strategie sperimentali per testare robustezza e generalizzazione:

1. Addestramento su singola coorte: Addestramento su dati di un solo lettore e test su entrambi i lettori (per valutare la sensibilità allo stile di annotazione).
2. Addestramento misto con Cross-Validation (5-fold): Utilizzo combinato dei dati di entrambi i lettori per valutare la generalizzazione inter-lettore.
3. Test basato sulla dimensione della ghiandola: Stratificazione dei dati in "grandi" e "piccole" ghiandole per valutare la robustezza dello sbilanciamento delle classi e della variabilità morfologica.

Ottimizzazione e Pre-processing:

• Pre-processing: Ricalibrazione a risoluzione isotropica ($0.5 \times 0.5 \times 3.0$ mm³), normalizzazione dell'intensità, e augmentation stocastica (zoom, crop, rotazioni, rumore).
• Ottimizzazione: I parametri iperparametrici (learning rate, dimensione del batch, dimensione delle feature) sono stati ottimizzati automaticamente utilizzando Optuna.
• Funzione di Loss: Una combinazione pesata di Dice Loss e Binary Cross-Entropy (BCE) per gestire lo sbilanciamento delle classi.
• Metrica: Il coefficiente di similarità di Dice (DSC) è stato utilizzato come metrica principale, con test statistici (Wilcoxon signed-rank e Friedman test) per la significatività.

3. Risultati Chiave

I risultati dimostrano che i modelli basati su Transformer superano la 3D U-Net, specialmente in scenari di variabilità e sbilanciamento.

• Addestramento su singola coorte:

  • SwinUNETR ha ottenuto i migliori risultati: DSC medio di 0.816 (Reader #1) e 0.860 (Reader #2).
  • UNETR ha ottenuto 0.800 e 0.833 rispettivamente.
  • La 3D U-Net (baseline) ha ottenuto 0.825 e 0.851.
  • Osservazione: I modelli Transformer hanno mostrato una maggiore stabilità quando testati su dati di un lettore diverso da quello di addestramento rispetto alla CNN.

• Addestramento Misto (Cross-Validation):

  • SwinUNETR ha raggiunto un DSC medio di 0.8583 (Reader #1) e 0.8678 (Reader #2).
  • UNETR ha ottenuto 0.8517 e 0.8640.
  • La 3D U-Net ha mostrato prestazioni inferiori (0.826 e 0.875).
  • L'analisi statistica (Friedman test) ha confermato differenze significative ($p < 0.0001$) tra le prestazioni dei modelli su Reader #1, con SwinUNETR in testa.

• Robustezza alla Dimensione della Ghiandola (Sbilanciamento Estremo):

  • Su dataset di ghiandole piccole, SwinUNETR ha mantenuto una stabilità eccezionale (DSC ~0.85) anche con proporzioni di dati ridotte, mentre UNETR e U-Net hanno mostrato cali significativi di prestazioni (fino a -0.044 per U-Net).
  • Su dataset di ghiandole grandi, SwinUNETR ha raggiunto un picco di 0.902 di DSC, superando nettamente gli altri modelli.
  • Le finestre spostate di SwinUNETR hanno mitigato efficacemente la sensibilità al rumore delle etichette e allo sbilanciamento delle classi.

4. Contributi Principali

1. Primo Benchmark Sistematico: Questo studio fornisce la prima analisi comparativa approfondita di UNETR e SwinUNETR su un dataset multi-lettore (546 volumi) per la segmentazione della prostata.
2. Mitigazione della Variabilità Inter-Lettore: Dimostrazione che i meccanismi di attenzione globale e a finestre spostate riducono la dipendenza dallo "stile" di annotazione specifico di un singolo radiologo, migliorando la generalizzazione.
3. Gestione dello Sbilanciamento delle Classi: Evidenza che le architetture Transformer, in particolare SwinUNETR, sono più robuste rispetto alle CNN tradizionali quando si affrontano ghiandole di piccole dimensioni o dataset con forte sbilanciamento foreground/background.
4. Ottimizzazione Automatizzata: Integrazione di Optuna per l'ottimizzazione degli iperparametri specifici per ogni architettura e dataset, garantendo un confronto equo.

5. Significato e Implicazioni

Il lavoro dimostra che l'uso di encoder basati su Transformer, in particolare SwinUNETR, rappresenta un passo avanti significativo per la segmentazione clinica della prostata.

• Affidabilità Clinica: La capacità di SwinUNETR di mantenere alte prestazioni indipendentemente dal lettore o dalla dimensione della ghiandola lo rende un candidato ideale per il deployment clinico, riducendo la necessità di ri-addestramento per ogni nuovo centro o radiologo.
• Efficienza Computazionale: Nonostante la complessità dei Transformer, SwinUNETR offre un ottimo compromesso tra accuratezza e efficienza computazionale, permettendo l'elaborazione in tempo quasi reale su GPU standard.
• Futuro della Segmentazione: I risultati suggeriscono che i modelli basati su attenzione globale possono superare i limiti delle CNN tradizionali nella medicina di precisione, offrendo soluzioni "reader-agnostic" (indipendenti dal lettore) che migliorano la coerenza nei processi diagnostici e di pianificazione terapeutica.