MSTNet: Diamo Nuova Vita alle Immagini con la Magia dei Transformer! (Anche in Medicina!)

Ciao a tutti! Oggi voglio parlarvi di qualcosa che mi appassiona tantissimo: come possiamo “guarire” le immagini rovinate, quelle sfocate o piene di rumore che a volte ci capitano tra le mani. Immaginate di avere una foto preziosa ma rovinata, o pensate a quanto sia cruciale avere immagini mediche perfette per una diagnosi accurata. Ecco, il restauro delle immagini è proprio questo: un processo affascinante e complesso per recuperare la chiarezza originale da un’immagine degradata.

Perché è così importante, soprattutto in campo medico?

Pensate alle endoscopie, come la colonscopia. È uno strumento fondamentale per scovare precocemente tumori o altre problematiche, ad esempio nel colon-retto (che, ahimè, è uno dei tumori più diffusi). La qualità delle immagini durante l’esame è tutto! Se l’immagine è sfocata per un movimento brusco, per problemi di messa a fuoco o per disturbi nel segnale, il medico potrebbe faticare a vedere piccoli polipi o lesioni sospette. Questo può compromettere la diagnosi e, di conseguenza, la salute del paziente.

Non solo, ma anche i sistemi di intelligenza artificiale che aiutano i medici (ad esempio per ricostruire in 3D il colon, localizzare polipi o classificarli) hanno bisogno di immagini impeccabili per funzionare al meglio. Se l’input è “sporco”, l’output non sarà affidabile. Quindi, migliorare la qualità delle immagini è il primo passo fondamentale.

Le vecchie strade: CNN e Transformer “classici”

Finora, per il restauro delle immagini si sono usate principalmente due tecnologie: le Reti Neurali Convoluzionali (CNN) e i modelli Transformer. Le CNN sono bravissime a catturare i dettagli locali, ma faticano un po’ a vedere il “quadro generale”, le relazioni a lunga distanza nell’immagine. I Transformer, nati nel mondo del linguaggio ma poi adattati alle immagini, sono invece campioni nel cogliere le dipendenze globali grazie a un meccanismo chiamato “auto-attenzione”. Però, a volte rischiano di concentrarsi troppo sul globale perdendo dettagli fini, o viceversa. Insomma, ognuno ha i suoi pro e contro.

La nostra idea: MSTNet, il meglio dei due mondi (e anche di più!)

Di fronte a queste sfide, ci siamo chiesti: “Possiamo creare qualcosa che unisca la forza del locale con quella del globale in modo intelligente?”. La risposta è sì, e l’abbiamo chiamata MSTNet (Multi-Stage Transformer Network).

L’idea di base è semplice ma potente:

• Usiamo un approccio multi-stadio progressivo: invece di fare tutto in un colpo solo, “ripuliamo” l’immagine passo dopo passo, affrontando il problema gradualmente.
• In ogni stadio, usiamo tipi diversi di Transformer per catturare informazioni differenti.
• Combiniamo Transformer basati su finestre (window-based), che sono ottimi per analizzare i dettagli all’interno di piccole aree locali, con Transformer che lavorano a livello di canali (channel-level), perfetti per capire le relazioni globali tra le caratteristiche dell’intera immagine.

È come avere due specialisti: uno che guarda i dettagli al microscopio e uno che osserva il panorama dall’alto, e poi farli collaborare!

Come funziona MSTNet nel dettaglio?

La nostra rete MSTNet è strutturata in tre stadi principali:

1. Primi due stadi: Focus sul Locale con Finestre. Qui usiamo una combinazione intelligente di una architettura U-Net (molto usata nel medical imaging) potenziata con i nostri blocchi Transformer basati su finestre (li chiamiamo WTB-UNet). Questi blocchi dividono l’immagine in piccole finestre e applicano l’attenzione solo all’interno di ciascuna finestra. Questo ci permette di catturare i dettagli spaziali locali in modo molto efficiente. Dopo ogni stadio, un modulo speciale (SAM) seleziona le informazioni più utili da passare allo stadio successivo.
2. Terzo stadio: Visione Globale sui Canali. Nei primi stadi, per analizzare l’immagine a diverse scale, si fanno operazioni di “compressione” (downsampling) che, inevitabilmente, possono far perdere qualche dettaglio fine. Per recuperare tutto al meglio, nel terzo stadio lavoriamo sull’immagine a risoluzione originale e usiamo i Transformer a livello di canali (CTB). Questi analizzano le relazioni tra i diversi canali di feature dell’immagine su tutta la sua estensione, catturando così il contesto globale e le informazioni semantiche.
3. La Fusione Fa la Forza: Abbiamo anche progettato un modulo di fusione speciale. Poiché le informazioni catturate dai Transformer a finestra e da quelli a canali sono diverse ma complementari, questo modulo le combina in modo dinamico e intelligente, creando una rappresentazione finale molto più ricca e accurata.

Per addestrare la rete, usiamo una combinazione di funzioni di costo (loss function) che tengono conto sia della precisione pixel per pixel (Charbonnier loss) sia della preservazione dei bordi e dei dettagli (Edge loss), soprattutto nell’ultimo stadio.

Mettiamolo alla prova: i risultati parlano!

Ovviamente, non basta avere una bella idea, bisogna dimostrare che funziona! Abbiamo testato MSTNet su diversi compiti classici di restauro immagini, come la rimozione della sfocatura (deblurring) e la rimozione del rumore (denoising).

Abbiamo usato dataset standard molto conosciuti (GoPro, HIDE, RealBlur per il deblurring; SIDD, DND per il denoising) e, cosa per noi molto importante, abbiamo creato e testato la nostra rete su un nostro dataset di immagini di colonscopia. Abbiamo preso video reali di colonscopie, selezionato frame nitidi e poi aggiunto artificialmente sfocatura (simulando problemi di focus) e rumore (simulando disturbi di segnale) per creare coppie di immagini degradate/originali.

I risultati? Beh, siamo davvero soddisfatti! Su quasi tutti i dataset e per entrambi i compiti, MSTNet ha superato i metodi allo stato dell’arte (SoTA), spesso con un margine significativo, sia in termini di metriche quantitative (PSNR e SSIM, che misurano la qualità e la somiglianza con l’originale) sia qualitativamente (le immagini restaurate sono visibilmente più nitide e pulite).

Ad esempio, nel deblurring sul dataset GoPro, abbiamo ottenuto un PSNR di 33.40 dB, superando il secondo miglior metodo di 0.34 dB. Anche sul nostro dataset del colon, MSTNet ha mostrato prestazioni eccellenti, raggiungendo 34.70 dB nel deblurring e 42.66 dB nel denoising, dimostrando la sua efficacia anche in questo specifico e delicato ambito medico.

Abbiamo anche fatto degli esperimenti “di ablazione”, cioè abbiamo provato a togliere o modificare parti della nostra rete (come il numero di stadi, il tipo di Transformer usato in ogni stadio, il modulo di fusione) per verificare che ogni componente contribuisse effettivamente al risultato finale. E sì, ogni scelta progettuale si è rivelata importante!

Non è tutto oro quel che luccica? Limiti e prospettive future

Siamo molto contenti dei risultati, ma siamo anche consapevoli che c’è sempre spazio per migliorare. Attualmente, MSTNet dà il meglio su degradazioni “semplici” (o sfocatura o rumore). Se un’immagine ha tanti problemi contemporaneamente (sfocatura, rumore, illuminazione scarsa…), la sfida diventa più complessa. Inoltre, per ora ci siamo concentrati su singole immagini. Estendere l’approccio ad altre modalità di imaging medico (come CT o MRI) o a video richiederà ulteriori adattamenti.

Per il futuro, vogliamo proprio lavorare su questo: rendere MSTNet ancora più robusto a degradazioni complesse, magari integrando conoscenze specifiche del dominio (ad esempio, informazioni sulla struttura del colon) e ottimizzando ulteriormente l’architettura.

In conclusione

MSTNet rappresenta un passo avanti significativo nel campo del restauro delle immagini. Combinando in modo innovativo Transformer locali e globali in un’architettura progressiva multi-stadio, siamo riusciti a ottenere risultati davvero notevoli, sia su benchmark generici che su dati medici specifici come quelli delle colonscopie.

Speriamo che il nostro lavoro possa non solo migliorare la qualità delle immagini in generale, ma anche, e soprattutto, aiutare i medici a fare diagnosi più precise e fornire dati migliori per tutti quei sistemi di AI che stanno rivoluzionando la medicina. La strada è ancora lunga, ma siamo convinti che tecnologie come MSTNet abbiano un potenziale enorme per “vedere” meglio e, in definitiva, prenderci cura meglio della nostra salute!

Fonte: Springer