SANA I2I: A Text Free Flow Matching Framework for Paired Image to Image Translation with a Case Study in Fetal MRI Artifact Reduction

1. Il Problema

La generazione di immagini ad alta risoluzione ha fatto passi da gigante grazie ai modelli di diffusione e Flow Matching. Tuttavia, architetture avanzate come SANA e la sua estensione SanaControlNet sono progettate principalmente per la generazione condizionata da testo (text-to-image) o da testo e immagine.
In molti scenari reali di traduzione immagine-immagine (I2I), come la rimozione di artefatti, il denoising o la traduzione di modalità, le annotazioni testuali sono inutili, ambigue o non disponibili. In questi casi, l'output desiderato è completamente specificato dall'immagine di input.
Il problema specifico affrontato nel paper è la riduzione degli artefatti da movimento nelle risonanze magnetiche (MRI) fetali. Ottenere dati reali appaiati (coppie di immagini "corrotte" e "pulite" dello stesso feto) è estremamente difficile a causa del movimento imprevedibile del feto e della madre. Inoltre, i metodi esistenti basati su apprendimento non appaiato spesso falliscono nel mantenere la fedeltà anatomica o non rimuovono efficacemente gli artefatti gravi.

2. Metodologia: SANA-I2I

Gli autori propongono SANA-I2I, un framework di traduzione immagine-immagine supervisionata, privo di testo, basato su Flow Matching nello spazio latente.

• Architettura e Rimozione del Testo: Il modello deriva da SanaControlNet ma rimuove completamente il condizionamento testuale. Invece di utilizzare un embedding di testo ($y$) e un segnale di controllo separato, SANA-I2I utilizza l'immagine di sorgente come unico segnale di condizionamento.
  • L'immagine di input (sorgente) viene elaborata attraverso lo stesso modulo di embedding (PatchEmbed) utilizzato per il rumore e il segnale di controllo, unificando i percorsi di elaborazione.
  • Il modello impara un campo di velocità condizionale che mappa la distribuzione del rumore verso la distribuzione dell'immagine target, guidato esclusivamente dall'immagine di sorgente.
• Meccanismo di Guida (CFG): Viene mantenuto il meccanismo di Classifier-Free Guidance (CFG). Durante l'addestramento, la variabile di condizionamento ($y$, ovvero l'immagine di sorgente) viene casualmente "droppata" (impostata a zero) con una certa probabilità. Questo permette di controllare la forza della trasformazione durante l'inferenza senza riaddestrare il modello.
• Varianti:
  • Approccio Primario: Solo il condizionamento $y$ viene droppato; il percorso di controllo rimane attivo. È ottimizzato per compiti di traduzione appaiata dove la fedeltà strutturale è cruciale.
  • Approccio "Bis" (Esplorativo): Sia $y$ che il segnale di controllo vengono droppati, permettendo una generazione incondizionata. Tuttavia, i risultati mostrano che questo approccio è inferiore per la rimozione degli artefatti fetali rispetto all'approccio primario.
• Efficienza: Il modello utilizza un Deep Compression Autoencoder (DC-AE) per comprimere le immagini in uno spazio latente compatto, riducendo drasticamente il numero di token da elaborare e permettendo la generazione ad alta risoluzione con costi computazionali ridotti.

3. Strategia dei Dati e Addestramento

Poiché i dati reali appaiati per le MRI fetali non sono disponibili, gli autori adottano una strategia di generazione di dati sintetici:

• Dataset: Utilizzano immagini MRI fetali reali pulite (Dominio B) e applicano un simulatore di artefatti basato sul metodo di Duffy et al. per creare immagini corrotte sintetiche (Dominio C).
• Configurazione: Vengono generati dati con artefatti di movimento severi (intervallo SSIM tra 0.6 e 0.9) per addestrare il modello a gestire degradazioni significative.
• Addestramento: Il modello è stato addestrato per 100 epoche con un batch size di 16, utilizzando l'ottimizzatore CAME e precisione mista (FP16). L'addestramento è stato condotto su immagini a 512x512 pixel.

4. Risultati Sperimentali

Lo studio valuta il metodo su due fronti: dati sintetici appaiati (C $\to$ B) e dati reali non appaiati (A $\to$ B, dove A sono immagini reali corrotte).

• Valutazione Quantitativa (Dati Sintetici):
  • Le metriche pixel-per-pixel come SSIM e MAE mostrano risultati leggermente inferiori per SANA-I2I rispetto al metodo precedente.
  • Spiegazione: Questo è dovuto al fatto che SANA-I2I corregge attivamente le inconsistenze di luminosità globale (che il metodo precedente lascia invariate). Sebbene questo aumenti l'errore pixel-per-pixel, produce immagini anatomicamente più corrette e prive di artefatti.
• Valutazione su Dati Reali (Metriche di Distribuzione):
  • Metriche basate sulla distribuzione come FID (Fréchet Inception Distance) e KID (Kernel Inception Distance) favoriscono chiaramente SANA-I2I.
  • Il modello proposto ottiene un FID di 51.74 contro 60.14 del metodo precedente, indicando che le immagini ricostruite sono statisticamente più vicine alla distribuzione delle MRI fetali pulite.
• Analisi Qualitativa:
  • SANA-I2I rimuove efficacemente gli artefatti da movimento (ghosting e distorsioni spaziali) preservando i bordi anatomici e i dettagli fini.
  • Il metodo precedente tende a lasciare residui di artefatti o a causare un eccessivo sfocamento (over-smoothing).
• Efficienza Computazionale:
  • Il modello raggiunge risultati di alta qualità con un numero molto ridotto di passi di inferenza (5 passi sono sufficienti, con risultati accettabili già a 2 passi).
  • Un'analisi di ablazione sul guidance scale mostra che valori tra 1.0 e 1.1 offrono il miglior compromesso tra rimozione degli artefatti e conservazione dei dettagli anatomici.

5. Contributi Chiave

1. Estensione Text-Free di SANA: Introduzione di un framework I2I supervisionato che rimuove la dipendenza dai prompt testuali, rendendo il modello adatto a compiti di trasformazione dove l'immagine di input definisce completamente l'output.
2. Allineamento al Task I2I: Dimostrazione che la rimozione del condizionamento testuale allinea meglio il modello ai compiti di trasformazione immagine-immagine, riducendo la stocasticità indesiderata e il numero di passi di inferenza necessari.
3. Efficienza e Scalabilità: Conferma che l'architettura basata su Flow Matching e DC-AE di SANA mantiene i vantaggi di efficienza e scalabilità anche in contesti puramente visivi ad alta risoluzione.
4. Validazione Clinica: Applicazione e validazione del metodo su un caso d'uso medico critico (riduzione artefatti MRI fetali), dimostrando la capacità di generalizzare su dati reali pur essendo addestrato su dati sintetici.

6. Significato e Impatto

Il lavoro di SANA-I2I è significativo perché colma il divario tra le potenti architetture di generazione basate su testo/immagine e le esigenze specifiche dei compiti di elaborazione medica supervisionata.

• Indipendenza dal Testo: Elimina la necessità di annotazioni testuali, spesso assenti o irrilevanti in ambito medico.
• Robustezza: Dimostra che i modelli generativi basati su Flow Matching possono essere addestrati su dati sintetici realistici per risolvere problemi complessi su dati reali, superando i limiti dei metodi non appaiati precedenti.
• Efficienza: La capacità di produrre risultati di alta qualità in pochi passi di inferenza rende il metodo promettente per applicazioni cliniche dove la velocità di elaborazione è fondamentale.

In sintesi, SANA-I2I rappresenta un passo avanti verso modelli generativi più specializzati ed efficienti per la ricostruzione e il miglioramento delle immagini mediche, offrendo una soluzione superiore per la correzione degli artefatti nelle MRI fetali rispetto alle tecniche esistenti.