NeuReg: Domain-invariant 3D Image Registration on Human and Mouse Brains

TL;DR

This work establishes a new state-of-the-art for domain-agnostic 3D brain image registration, underpinned by Neuro-inspired Transformer-based architecture that generates domain-agnostic representations of imaging features and incorporates a shifting window-based Swin Transformer block as the encoder.

Abstract

Medical brain imaging relies heavily on image registration to accurately curate structural boundaries of brain features for various healthcare applications. Deep learning models have shown remarkable performance in image registration in recent years. Still, they often struggle to handle the diversity of 3D brain volumes, challenged by their structural and contrastive variations and their imaging domains. In this work, we present NeuReg, a Neuro-inspired 3D image registration architecture with the feature of domain invariance. NeuReg generates domain-agnostic representations of imaging features and incorporates a shifting window-based Swin Transformer block as the encoder. This enables our model to capture the variations across brain imaging modalities and species. We demonstrate a new benchmark in multi-domain publicly available datasets comprising human and mouse 3D brain volumes. Extensive experiments reveal that our model (NeuReg) outperforms the existing baseline deep learning-based image registration models and provides a high-performance boost on cross-domain datasets, where models are trained on 'source-only' domain and tested on completely 'unseen' target domains. Our work establishes a new state-of-the-art for domain-agnostic 3D brain image registration, underpinned by Neuro-inspired Transformer-based architecture.

Figures (5)

• Figure 1: Block diagram for the proposed methodology. Input (fixed and moving) samples are passed to the domain generalised layer followed by patch partition which splits the input samples into patches which are fed to the swin transformer encoder referred to as Stage 1. Our architecture consists of four such repeated stages after which the the Swin transformer encoder generates the deformation field. A Discrete Fourier Transform (DFT) is applied to convert the deformation field to a fourier domain, which is then passed through the model-driven decoder. The moving brain segmentation is warped with the deformation field via a spatial transform to generate the final output.
• Figure 1: Qualitative results on iSeg-2017 dataset. Top rows show the zoomed-in version of the samples shown in the bottom row with DICE score mentioned underneath. We show the segmentation results of our model with NCC and MSE along with SynthMorph and FourierNet.
• Figure 2: Qualitative results on iSeg-2017 dataset. Top rows show the zoomed-in version of the normal and our model's data representation, followed by the ground truth segmentation, aligned moving segmentation generated by our model with NCC and MSE and the aligned moving segmentation generated by SynthMorph and FourierNet. Bottom rows show the corresponding samples with the DICE score mentioned underneath.
• Figure 2: Qualitative result on DevCCF dataset's P04, P14 and P56 samples. The top row shows the original samples for each domain. The middle row shows the registration result of our model whereas the last row is the FourierNet's results. The DICE score between the aligned moving brain segmentation and the fixed segmentation is mentioned at the bottom right of each sample.
• Figure 3: Qualitative results on DevCCF dataset. The original samples from all the domains are shown (top row) followed by the output of our top-performing model (middle row) and FourierNet output (bottom row) with the corresponding SSIM mentioned underneath.