Source-linked AI summary
DeepAtlas: Joint Semi-Supervised Learning of Image Registration and Segmentation
Zhenlin Xu, Marc Niethammer
TL;DR
Manual 3D medical-image segmentations are difficult and labor intensive to obtain, limiting supervised segmentation when labels are scarce. DeepAtlas jointly learns registration and segmentation, using existing or predicted segmentations for registration supervision and registration-based augmentation for segmentation. Across knee and brain MRIs, it improves both tasks over separately trained networks, including one-shot Dice gains over unsupervised registration of 2.7% and 1.8%.
Problem
Obtaining the many manual 3D segmentations typically needed for supervised CNN training is difficult and labor intensive.
Method
DeepAtlas jointly trains weakly supervised registration and semi-supervised segmentation networks that mutually guide each other on unlabeled images through anatomy similarity.
Results
DeepAtlas produces large simultaneous improvements in segmentation and registration accuracy over separately trained networks on knee and brain MRIs, including one-shot Dice gains over unsupervised registration of 2.7% on knee and 1.8% on brain data.
Takeaways & Limitations
The framework supports high-quality joint models with very limited labeled data, including one-shot segmentation learning and improved registration.
Abstract
from arXiv · showhide
Deep convolutional neural networks (CNNs) are state-of-the-art for semantic image segmentation, but typically require many labeled training samples. Obtaining 3D segmentations of medical images for supervised training is difficult and labor intensive. Motivated by classical approaches for joint segmentation and registration we therefore propose a deep learning framework that jointly learns networks for image registration and image segmentation. In contrast to previous work on deep unsupervised image registration, which showed the benefit of weak supervision via image segmentations, our approach can use existing segmentations when available and computes them via the segmentation network otherwise, thereby providing the same registration benefit. Conversely, segmentation network training benefits from the registration, which essentially provides a realistic form of data augmentation. Experiments on knee and brain 3D magnetic resonance (MR) images show that our approach achieves large simultaneous improvements of segmentation and registration accuracy (over independently trained networks) and allows training high-quality models with very limited training data. Specifically, in a one-shot-scenario (with only one manually labeled image) our approach increases Dice scores (%) over an unsupervised registration network by 2.7 and 1.8 on the knee and brain images respectively.
1 Introduction
Medical image segmentation and registration are interrelated, but deep learning methods face limited availability of manual 3D segmentations. DeepAtlas addresses this gap by jointly learning registration and segmentation networks using estimated supervision from unlabeled images.
- Motivation: Segmentation and registration can provide complementary information, motivating approaches that jointly estimate both tasks.Segmentations can supervise registration and evaluate its results, while registered labeled atlases can support segmentation.
- Motivation: CNN segmentation and deep registration are effective with sufficient supervision, but obtaining manual 3D medical-image segmentations is difficult and labor intensive.Registration training may use optimization-based estimates or unsupervised image-only objectives, while manual segmentations are often unavailable for much image data.
- Motivation: Weak supervision from image segmentations can improve registration over unsupervised training, creating a need to exploit segmentation information when manual labels are scarce.The additional segmentation loss improves registration, but large datasets often lack manual segmentations.
- Proposed approach: DeepAtlas jointly learns deep networks for weakly supervised registration and semi-supervised segmentation.The framework is designed to use few manual segmentations while learning both tasks together.
- Proposed approach: Unlike prior joint methods that optimize each image pair, DeepAtlas learns from an image population and allows the resulting networks to operate independently at test time.Its networks mutually guide training on unlabeled images through an anatomy similarity loss that compares warped and target segmentations.
2 Method
DeepAtlas jointly trains registration and segmentation networks, using anatomy-based weak supervision to connect their learning when manual segmentations are scarce. The method alternates network optimization and supplies estimated segmentations for unlabeled registration pairs.
- Weakly-supervised Registration Learning: Registration optimizes intensity similarity, deformation regularization, and anatomy similarity between warped moving and target segmentations.The anatomy loss adds segmentation-based weak supervision to image-based registration.
- Weakly-supervised Registration Learning: The registration network predicts a displacement field that warps the moving image into target-image space.The deformation map is Φ^-1 = u + id, where id is the identity transform.
- Weakly-supervised Registration Learning: When manual segmentations are missing, the segmentation network estimates them so every training image pair can receive anatomy-based weak supervision.This extends registration supervision beyond pairs with existing manual labels.
- Semi-supervised Segmentation Learning: The segmentation network predicts probabilistic maps and learns from both manual segmentation loss and anatomy similarity when predictions participate in registration.For unlabeled images, registered atlas segmentations provide noisy labels through the anatomy loss.
- Implementation Details: The framework uses soft multi-class Dice for anatomy and supervised segmentation losses, normalized cross correlation for intensity similarity, and bending energy for regularization.Bending energy is a second-order generalization of diffusion regularization based on displacement-field Hessians.
- Implementation Details: Training alternates between the two networks because whole 3D images exceed GPU memory for simultaneous optimization.The implementation uses a 1:20 segmentation-to-registration training-step ratio and independently pretrains both networks.
3 Experiments and Results
Experiments on 3D knee and brain MRIs evaluate DeepAtlas against separately trained networks under limited manual supervision. Joint training improves segmentation and registration, including in one-shot learning.
- Experimental setup: DeepAtlas is evaluated on 3D knee and brain MRI datasets with N of M images manually labeled, where N << M.The knee dataset uses 200 training images, while the brain dataset uses 65 training images.
- Experimental setup: DA models are initialized from mono-networks, while one-shot DA training starts from Semi-DA segmentation and unsupervised registration networks.Semi-DA keeps one pretrained network fixed; DA alternately trains both networks.
- General results: DA and Semi-DA achieve significantly larger Dice scores than mono-networks across all experiments, with DA usually performing best.The comparison uses Mann-Whitney U-tests with multiple-comparison correction.
- Knee results: On knee MRIs, segmentation improves by about 1.2 and 0.5, while registration improves by about 3.1 and 3.0 with 5 and 10 manual segmentations.For cartilage structures, the corresponding gains are 1.4 and 0.7 for segmentation and 5.5 and 5.2 for registration.
- One-shot and qualitative results: 2.7 and 1.8 Dice-score gains over unsupervised registration are achieved on knee and brain data, respectively, in one-shot experiments.The qualitative results also show more anatomically consistent knee registrations for DA than mono-networks.
4 Conclusion
The conclusion presents DeepAtlas as a joint framework for learning segmentation and registration with few manual segmentations. It reports improvements across knee and brain MRI tasks and strong one-shot registration gains.
- DeepAtlas jointly learns segmentation and registration networks using only a few images with manual segmentations.
- An anatomical similarity loss makes learned registrations more anatomically consistent, while registration provides data augmentation for segmentation on unlabeled images.
- The framework improves bone, cartilage, and cortical-structure results over separately learned networks and supports one-shot segmentation learning.
A Supplementary material
The supplementary figures show the segmentation and registration architectures and provide qualitative brain registration and knee segmentation examples across models.
- Figure 3 contrasts the segmentation-network architecture with the registration-network architecture.Max-pooling performs down-sampling in segmentation, whereas 2-stride convolution performs it in registration.
- Figure 4 presents brain MRI registration examples in the top panel and knee MRI segmentation examples in the bottom panel.
- Figure 4 compares Mono-i and DA-i models trained with i manual segmentations alongside moving, target, original, and manual-reference images.