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

1. Probleemstelling

Niet-lineaire beeldregistratie is cruciaal voor medische beeldvorming, met name in de computationele neuroanatomie voor het uitlijnen van hersenscans. Hoewel Deep Learning (DL) dit veld heeft revolutionair door inferentie bijna direct mogelijk te maken, lijden bestaande DL-gebaseerde methoden aan aanzienlijke beperkingen:

• Hoge Rekenkosten: De meeste methoden vertrouwen op dichte 3D-roosterrepresentaties voor invoerbeelden, feature maps en vervormingsvelden, wat leidt tot massaal geheugengebruik (vaak >24GB RAM vereisend) en zware hardware-eisen.
• Gebrek aan Diffeomorfisme-Garanties in Oppervlakmethodes: Hoewel er light-weight methoden bestaan die grensvlakoppervlakken van gesegmenteerde regio's gebruiken om complexiteit te verminderen, missen deze doorgaans een theoretisch kader om meerdere regionale mappingen te fuseren tot één enkele, globaal consistente diffeomorfische transformatie (een gladde, inverteerbare mapping die topologie behoudt).

De auteurs beogen een light-weight, snel en diffeomorfisch DL-registratiekader te ontwikkelen dat oppervlakrepresentaties van anatomische structuren benut in plaats van dichte volumetrische roosters.

2. Methodologie: NimbleReg

NimbleReg herformuleert beeldregistratie als de uitlijning van grensvlakoppervlakken van meerdere gesegmenteerde anatomische regio's. Het kader bestaat uit drie hoofdfasen:

A. Preprocessing

1. Oppervlakte-extractie: Met behulp van tools zoals Marching Cubes worden oppervlakken geëxtraheerd uit gesegmenteerde regio's. Cruciaal is dat gedeelde interfaces tussen aangrenzende regio's worden verwerkt om identieke puntcoördinaten op gedeelde grenzen te garanderen.
2. Mesh-verwerking: Individuele oppervlakken worden samengevoegd tot één mesh, gladgestreken om interfaces te behouden, en vervolgens opnieuw opgesplitst in individuele anatomische oppervlakken.
3. Pre-uitlijning: Een polyaffiene transformatie (met behulp van Polaffini) lijnt het bewegende beeld uit op een standaardreferentie op basis van regio-centroïden. Dit verkleint de zoekruimte voor het DL-model en elimineert de noodzaak om globale pose of laagfrequente verschuivingen te leren.

B. Snelheidsschatting (Het Deep Learning-model)

• Architectuur: Het kernmodel, $f_\theta$, maakt gebruik van een PointNet-ruggengraat. Het verwerkt puntwolken (bewegende en referentie-oppervlakken) met een gedeelde Multi-Layer Perceptron (MLP) om lokale features te extraheren, die via max-pooling worden geaggregeerd tot permutatie-invariante globale features.
• Output: In plaats van punten direct te vervormen, geeft het model snelheidsvectoren ($v$) uit die geassocieerd zijn met de bewegende punten. Dit behandelt de output als een sparse parametrisatie van een snelheidsveld in plaats van een direct verplaatsingsveld.
• Training: Het model wordt getraind op een zelftoezicht (self-supervised) manier.
  • Fitting-verlies ($L_{fit}$): Gebruikt de symmetrische Chamfer-afstand om de kloof tussen getransformeerde bewegende punten en referentiepunten te minimaliseren.
  • Regularisatie-verlies ($L_{reg}$): Straft abrupte snelheidsveranderingen tussen punten binnen dezelfde simplex (driehoekig vlak) van het oppervlakmesh af om gladheid langs het oppervlak te garanderen.

C. Integratie in Diffeomorfische Transformatie

Dit is de belangrijkste innovatie die NimbleReg onderscheidt van eerdere oppervlakgebaseerde methoden.

1. Stationair Snelheidsveld (SVF) Constructie: De door $f_\theta$ gegenereerde sparse snelheden worden omgezet in een glad, continu SVF ($V$) met behulp van kernelconvolutie.
2. Log-Euclidisch Kader: Het SVF wordt geïntegreerd over tijd met gebruikmaking van de Lie-groepstructuur van diffeomorfismen. De uiteindelijke transformatie $\phi$ wordt berekend als de Lie-exponentiële afbeelding: $\phi = \exp(V)$.
3. Fusie: Omdat het SVF-kader lineaire operaties in de Lie-algebra toestaat, kunnen snelheden uit meerdere anatomische regio's naadloos worden gemengd. De uiteindelijke transformatie is de compositie van de pre-uitlijning, de geleerde diffeomorfie en de inverse pre-uitlijning.
4. Inferentie: De resulterende transformatie kan worden toegepast op elk vraagpunt (bijv. een voxelrooster) of beeld via achterwaartse coördinatenmapping.

3. Belangrijkste Bijdragen

• Light-weight Architectuur: Door dichte 3D-roosters te vervangen door PointNet-gebaseerde oppervlakverwerking, verkleint het model de geheugenvoetafdruk en de rekenlast drastisch.
• Diffeomorfisme-Garantie via SVF: Het artikel overbrugt succesvol de kloof tussen oppervlakgebaseerde registratie en globaal diffeomorfisme door het Log-Euclidisch kader en Stationaire Snelheidsvelden te integreren. Dit garandeert dat de uiteindelijke transformatie glad, inverteerbaar en topologiebehoudend is.
• Multi-Regio Fusie: De methode biedt een mechanisme om onafhankelijke regionale uitlijningen te fuseren tot één enkele, coherente globale transformatie, waarmee een beperking van eerdere oppervlakgebaseerde DL-benaderingen wordt opgelost.
• Efficiëntie: De aanpak maakt snelle inferentie mogelijk met minimaal geheugengebruik in vergelijking met state-of-the-art roostergebaseerde DL-methoden.

4. Experimentele Resultaten

De methode werd geëvalueerd op een niet-lineaire hersenregistratietaken met T1-gewogen MRI-gegevens van de ADNI, UK Biobank en IXI datasets (95 anatomische regio's).

• Vergelijking: NimbleReg werd vergeleken met EasyReg (een state-of-the-art DL-methode gebaseerd op SynthMorph) en traditionele affiene/polyaffiene baselines.
• Nauwkeurigheidsmetrieken:
  • Jaccard-index (Volumegedeelte): NimbleReg behaalde prestaties vergelijkbaar met EasyReg in corticale, subcorticale en ventriculaire regio's.
  • Chamfer-afstand (Oppervlakte-uitlijning): NimbleReg presteerde beter dan EasyReg in oppervlakgebaseerde metrieken. De auteurs merken op dat dit deels komt omdat NimbleReg expliciet is getraind op oppervlakte-uitlijning, terwijl EasyReg dat niet is.
• Rekenefficiëntie:
  • Geheugengebruik: EasyReg vereist ongeveer 24 GB RAM voor een batchgrootte van 1 vanwege zware CNN-feature maps. NimbleReg vereist <1 GB voor de snelheidsschattingfase, waarbij de integratiestap het enige geheugenintensieve deel is (gemitigeerd door KD-bomen).
  • Snelheid: De PointNet-ruggengraat maakt uiterst snelle snelheidsschatting mogelijk.

5. Betekenis en Toekomstig Werk

Betekenis:
NimbleReg demonstreert dat hoogwaardige, diffeomorfische beeldregistratie geen zware, roostergebaseerde deep learning-modellen vereist. Door gebruik te maken van de geometrische eigenschappen van gesegmenteerde oppervlakken en de wiskundige striktheid van het Log-Euclidisch kader, biedt het een schaalbaar, geheugenefficiënt alternatief dat geschikt is voor omgevingen met beperkte middelen of grootschalige studies.

Beperkingen & Toekomstige Richtingen:

• Corticale Resolutie: De nauwkeurigheid in corticale regio's kan worden verbeterd met hogere puntendichtheden, aangezien huidige decimatie fijne details kan verliezen.
• Kernelbandbreedte: De vaste bandbreedte van de integratiekernel kan leiden tot overmatig gladde transformaties in sparse gebieden; een adaptieve $\sigma$ gebaseerd op puntendichtheid wordt voorgesteld.
• Integratieschema: Het huidige eerste-orde Runge-Kutta (Euler) integratieschema kan worden opgewaardeerd naar methoden van hogere orde voor betere numerieke precisie.

Kortom, NimbleReg vertegenwoordigt een belangrijke stap naar praktische, efficiënte en wiskundig robuuste deep learning-oplossingen voor medische beeldregistratie.