NimbleReg: A light-weight deep-learning framework for diffeomorphic image registration

1. Enunciato del Problema

La registrazione non lineare delle immagini è fondamentale per l'imaging medico, in particolare nella neuroanatomia computazionale per l'allineamento delle scansioni cerebrali. Sebbene l'Apprendimento Profondo (DL) abbia rivoluzionato questo campo consentendo inferenze quasi istantanee, i metodi esistenti basati sul DL soffrono di limitazioni significative:

• Elevato Costo Computazionale: La maggior parte dei metodi si basa su rappresentazioni a griglia 3D dense per immagini in input, mappe di caratteristiche e campi di deformazione, portando a un enorme consumo di memoria (spesso richiedente >24 GB di RAM) e a requisiti hardware pesanti.
• Mancanza di Garanzie Diffeomorfe nei Metodi Basati su Superfici: Sebbene esistano metodi leggeri che utilizzano le superfici di confine delle regioni segmentate per ridurre la complessità, tipicamente mancano di un quadro teorico per fondere mappature regionali multiple in un'unica trasformazione diffeomorfa globalmente coerente (una mappatura liscia e invertibile che preserva la topologia).

Gli autori mirano a sviluppare un framework di registrazione DL leggero, veloce e diffeomorfo che sfrutti le rappresentazioni superficiali delle strutture anatomiche anziché griglie volumetriche dense.

2. Metodologia: NimbleReg

NimbleReg riformula la registrazione delle immagini come l'allineamento delle superfici di confine provenienti da multiple regioni anatomiche segmentate. Il framework è composto da tre fasi principali:

A. Preprocessing

1. Estrazione delle Superfici: Utilizzando strumenti come Marching Cubes, le superfici vengono estratte dalle regioni segmentate. Crucialmente, le interfacce condivise tra regioni adiacenti vengono elaborate per garantire coordinate puntuali identiche sui confini condivisi.
2. Elaborazione della Mesh: Le singole superfici vengono fuse in una mesh unica, levigate per preservare le interfacce e successivamente ridivise in singole superfici anatomiche.
3. Pre-allineamento: Una trasformazione poliaffine (utilizzando Polaffini) allinea l'immagine in movimento a un riferimento standard basato sui centroidi delle regioni. Questo riduce lo spazio di ricerca per il modello DL, eliminando la necessità di apprendere la posa globale o spostamenti a bassa frequenza.

B. Stima della Velocità (Il Modello di Deep Learning)

• Architettura: Il modello centrale, $f_\theta$, utilizza un backbone PointNet. Elabora nuvole di punti (superfici in movimento e di riferimento) utilizzando un Perceptron Multistrato (MLP) condiviso per estrarre caratteristiche locali, che vengono aggregate tramite max-pooling per creare caratteristiche globali invarianti alla permutazione.
• Output: Invece di deformare direttamente i punti, il modello restituisce vettori di velocità ($v$) associati ai punti in movimento. Questo tratta l'output come una parametrizzazione sparsa di un campo di velocità anziché un campo di spostamento diretto.
• Addestramento: Il modello è addestrato in modo auto-supervisionato.
  • Loss di Adattamento ($L_{fit}$): Utilizza la distanza Chamfer simmetrica per minimizzare il divario tra i punti in movimento trasformati e i punti di riferimento.
  • Loss di Regularizzazione ($L_{reg}$): Penalizza i cambiamenti bruschi di velocità tra punti all'interno dello stesso simplex (faccia triangolare) della mesh superficiale per garantire la levigatezza lungo la superficie.

C. Integrazione nella Trasformazione Diffeomorfa

Questa è la novità chiave che distingue NimbleReg dai precedenti metodi basati su superfici.

1. Costruzione del Campo di Velocità Stazionario (SVF): Le velocità sparse output da $f_\theta$ vengono convertite in un SVF ($V$) liscio e continuo utilizzando la convoluzione a kernel.
2. Quadro Log-Euclideo: L'SVF viene integrato nel tempo utilizzando la struttura del gruppo di Lie delle diffeomorfie. La trasformazione finale $\phi$ è calcolata come mappa esponenziale di Lie: $\phi = \exp(V)$.
3. Fusione: Poiché il quadro SVF consente operazioni lineari nell'algebra di Lie, le velocità provenienti da multiple regioni anatomiche possono essere fuse senza soluzione di continuità. La trasformazione finale è la composizione del pre-allineamento, della diffeomorfia appresa e del pre-allineamento inverso.
4. Inferenza: La trasformazione risultante può essere applicata a qualsiasi punto di interrogazione (ad esempio, una griglia di voxel) o immagine tramite mappatura inversa delle coordinate.

3. Contributi Chiave

• Architettura Leggera: Sostituendo le griglie 3D dense con l'elaborazione superficiale basata su PointNet, il modello riduce drasticamente l'impronta di memoria e il carico computazionale.
• Garanzia Diffeomorfa tramite SVF: Il paper colma con successo il divario tra la registrazione basata su superfici e la diffeomorfia globale integrando il quadro Log-Euclideo e i Campi di Velocità Stazionari. Questo garantisce che la trasformazione finale sia liscia, invertibile e preservante la topologia.
• Fusione Multi-Regionale: Il metodo fornisce un meccanismo per fondere allineamenti regionali indipendenti in un'unica trasformazione globale coerente, risolvendo una limitazione dei precedenti approcci DL basati su superfici.
• Efficienza: L'approccio consente un'inferenza rapida con un utilizzo minimo di memoria rispetto ai metodi DL basati su griglia all'avanguardia.

4. Risultati Sperimentali

Il metodo è stato valutato su un compito di registrazione cerebrale non lineare utilizzando dati MRI pesati in T1 provenienti dai dataset ADNI, UK Biobank e IXI (95 regioni anatomiche).

• Confronto: NimbleReg è stato confrontato con EasyReg (un metodo DL all'avanguardia basato su SynthMorph) e con baseline tradizionali affini/poliaffini.
• Metriche di Accuratezza:
  • Indice di Jaccard (Sovrapposizione Volumetrica): NimbleReg ha ottenuto prestazioni comparabili a EasyReg attraverso regioni corticali, sub-corticali e ventricolari.
  • Distanza Chamfer (Allineamento Superficiale): NimbleReg ha superato EasyReg nelle metriche basate su superficie. Gli autori notano che ciò è parzialmente dovuto al fatto che NimbleReg è esplicitamente addestrato sull'allineamento superficiale, mentre EasyReg non lo è.
• Efficienza Computazionale:
  • Utilizzo di Memoria: EasyReg richiede circa 24 GB di RAM per una dimensione del batch di 1 a causa delle pesanti mappe di caratteristiche CNN. NimbleReg richiede <1 GB per la fase di stima della velocità, con il passaggio di integrazione che è l'unica parte ad alta intensità di memoria (mitigata dagli alberi KD).
  • Velocità: Il backbone PointNet consente una stima della velocità estremamente rapida.

5. Significato e Lavori Futuri

Significato:
NimbleReg dimostra che una registrazione di immagini diffeomorfa di alta qualità non richiede pesanti modelli di deep learning basati su griglia. Sfruttando le proprietà geometriche delle superfici segmentate e il rigore matematico del quadro Log-Euclideo, offre un'alternativa scalabile ed efficiente dal punto di vista della memoria, adatta ad ambienti con risorse limitate o a studi su larga scala.

Limitazioni e Direzioni Future:

• Risoluzione Corticale: L'accuratezza nelle regioni corticali potrebbe essere migliorata con densità puntuali più elevate, poiché la decimazione attuale potrebbe perdere dettagli fini.
• Larghezza di Banda del Kernel: La larghezza di banda fissa del kernel di integrazione può portare a trasformazioni eccessivamente lisce in aree sparse; è suggerito un $\sigma$ adattivo basato sulla densità dei punti.
• Schema di Integrazione: L'attuale integrazione di Runge-Kutta del primo ordine (Eulero) potrebbe essere aggiornata a metodi di ordine superiore per una migliore precisione numerica.

In conclusione, NimbleReg rappresenta un passo significativo verso soluzioni pratiche, efficienti e matematicamente robuste di deep learning per la registrazione di immagini mediche.