Semi-weakly-supervised neural network training for medical image registration

TL;DR

This work tackles the challenge of deformable medical image registration when voxel-level supervision is scarce by proposing a semi-weakly-supervised framework that uses a small ROI-labelled set together with a large unlabelled dataset. It introduces two commutative image perturbations, WarpDDF and RegCut, along with a weight-perturbation consistency regime to enforce stable predictions on unlabelled data. The approach yields significant accuracy gains on a pelvic MRI dataset with eight ROIs and enables construction of a pelvic atlas, illustrating practical clinical utility while reducing annotation effort. Overall, the method combines ROI-guided supervision with self-supervised consistency and spatial augmentations to advance registration performance in low-label regimes.

Abstract

For training registration networks, weak supervision from segmented corresponding regions-of-interest (ROIs) have been proven effective for (a) supplementing unsupervised methods, and (b) being used independently in registration tasks in which unsupervised losses are unavailable or ineffective. This correspondence-informing supervision entails cost in annotation that requires significant specialised effort. This paper describes a semi-weakly-supervised registration pipeline that improves the model performance, when only a small corresponding-ROI-labelled dataset is available, by exploiting unlabelled image pairs. We examine two types of augmentation methods by perturbation on network weights and image resampling, such that consistency-based unsupervised losses can be applied on unlabelled data. The novel WarpDDF and RegCut approaches are proposed to allow commutative perturbation between an image pair and the predicted spatial transformation (i.e. respective input and output of registration networks), distinct from existing perturbation methods for classification or segmentation. Experiments using 589 male pelvic MR images, labelled with eight anatomical ROIs, show the improvement in registration performance and the ablated contributions from the individual strategies. Furthermore, this study attempts to construct one of the first computational atlases for pelvic structures, enabled by registering inter-subject MRs, and quantifies the significant differences due to the proposed semi-weak supervision with a discussion on the potential clinical use of example atlas-derived statistics.

Figures (4)

• Figure 1: The proposed semi-weakly-supervised registration pipeline with $L_{weakly\_supervised}$ calculated from the labelled pair and $L_\text{cons}$ calculated from the unlabelled pair.
• Figure 2: DDF-based augmentation transforms the fixed image $I^\text{fix}_{(\text{unl})}$ by an augmentation dense displacement field $\mathbb{u}^\text{aug}$ and seek consistency between DDFs predicted from the original and augmented pairs. Note segmentation masks are not available during training, they are included here for illustration purpose.
• Figure 3: RegCut mix unlabelled moving image with unlabelled fixed image following a randomly sampled mask and seek consistency between DDFs predicted from the original and augmented pairs. Note segmentation masks are not available during training, they are included here for illustration purpose.
• Figure 4: (a) The resulting atlas and associated segmentation masks. (b) The probability maska of the atlas for the eight structures.