Robust Rigid and Non-Rigid Medical Image Registration Using Learnable Edge Kernels

TL;DR

The paper tackles the challenge of robust multi-modal medical image registration by introducing learnable edge kernels that bias feature extraction toward anatomical boundaries. Integrated into separate rigid (affine) and non-rigid (deformable) networks, these kernels are initialized with edge priors and refined during unsupervised training, enabling more accurate alignment of complex brain structures across modalities. Eight model variants (four rigid, four non-rigid) are systematically evaluated on a mix of in-house and public 3D MRI datasets, with Monte Carlo dropout used to quantify uncertainty. Across rigid and non-rigid scenarios, the edge-guided architectures consistently outperform established baselines, highlighting the value of edge-focused representations for enhancing multi-modal and structural registration in neuroimaging.

Abstract

Medical image registration is crucial for various clinical and research applications including disease diagnosis or treatment planning which require alignment of images from different modalities, time points, or subjects. Traditional registration techniques often struggle with challenges such as contrast differences, spatial distortions, and modality-specific variations. To address these limitations, we propose a method that integrates learnable edge kernels with learning-based rigid and non-rigid registration techniques. Unlike conventional layers that learn all features without specific bias, our approach begins with a predefined edge detection kernel, which is then perturbed with random noise. These kernels are learned during training to extract optimal edge features tailored to the task. This adaptive edge detection enhances the registration process by capturing diverse structural features critical in medical imaging. To provide clearer insight into the contribution of each component in our design, we introduce four variant models for rigid registration and four variant models for non-rigid registration. We evaluated our approach using a dataset provided by the Medical University across three setups: rigid registration without skull removal, with skull removal, and non-rigid registration. Additionally, we assessed performance on two publicly available datasets. Across all experiments, our method consistently outperformed state-of-the-art techniques, demonstrating its potential to improve multi-modal image alignment and anatomical structure analysis.

Figures (21)

• Figure 1: Top: Architecture of the learnable edge detection module. Bottom: Features extracted by the learnable edge detection module.
• Figure 2: Residual Block and Dilated Convolution Block Architectures: The residual block (top) uses convolutional layers with shortcut connections for efficient gradient flow, while the dilated convolution block (bottom) employs multiple dilations for multi-scale contextual information capture.
• Figure 3: The inception module consists of parallel convolutional layers with different filter sizes ($1 \times 1$, $3 \times 3$, and $5 \times 5$), allowing for multi-scale feature extraction. The outputs of these layers are concatenated along the channel dimension.
• Figure 4: The architecture of the Dense Feature Fusion Module.
• Figure 5: Reg-LEdge-Model Variant-1 architecture (rigid): takes Laplacian convoluted moving and fixed volumes and outputs affine parameters