Source-linked AI summary
Data augmentation using learned transformations for one-shot medical image segmentation
Amy Zhao, Guha Balakrishnan, Frédo Durand, John V. Guttag, Adrian V. Dalca
TL;DR
Limited labeled medical images and imperfect hand-tuned augmentation motivate a method that learns realistic transformations from unlabeled scans. The method applies learned spatial and appearance transformations to one labeled scan, and the resulting examples improve one-shot brain MRI segmentation, approaching fully supervised performance. The demonstrated scope is brain MRI, with extensions to other anatomies and modalities left feasible rather than evaluated.
Problem
Medical segmentation needs large labeled datasets for supervised deep learning, but labeling requires substantial expertise and time and hand-tuned augmentation misses complex image variation.
Method
The method learns independent spatial and appearance transformation models from unlabeled images, samples them, and applies them to a single labeled scan to synthesize labeled examples.
Results
The synthesized examples train a supervised segmenter that outperforms existing one-shot methods on every test example and approaches fully supervised performance.
Takeaways & Limitations
Learned transformation-based augmentation enables one-shot segmentation using a single labeled scan and unlabeled images, including in settings where only a few scans can be annotated.
Takeaways & Limitations
The approach was demonstrated on brain MRI, although the authors state that its lack of brain- or MRI-specific information makes extension to other anatomy or modalities feasible.
Abstract
from arXiv · showhide
Image segmentation is an important task in many medical applications. Methods based on convolutional neural networks attain state-of-the-art accuracy; however, they typically rely on supervised training with large labeled datasets. Labeling medical images requires significant expertise and time, and typical hand-tuned approaches for data augmentation fail to capture the complex variations in such images. We present an automated data augmentation method for synthesizing labeled medical images. We demonstrate our method on the task of segmenting magnetic resonance imaging (MRI) brain scans. Our method requires only a single segmented scan, and leverages other unlabeled scans in a semi-supervised approach. We learn a model of transformations from the images, and use the model along with the labeled example to synthesize additional labeled examples. Each transformation is comprised of a spatial deformation field and an intensity change, enabling the synthesis of complex effects such as variations in anatomy and image acquisition procedures. We show that training a supervised segmenter with these new examples provides significant improvements over state-of-the-art methods for one-shot biomedical image segmentation. Our code is available at https://github.com/xamyzhao/brainstorm.
1. Introduction
Biomedical image segmentation supports important clinical and research applications, but supervised deep learning requires labeled data that is costly and time-consuming to obtain.
- Semantic segmentation supports population analyses, disease diagnosis, and treatment planning.
- Supervised deep learning achieves state-of-the-art segmentation when enough labeled data is available.
- Manual segmentation labels require considerable expertise and time.
MIT, MGH
Medical images vary substantially across anatomy, contrast, texture, and acquisition conditions, while hand-tuned augmentation has limited ability to reproduce these variations. The proposed method learns transformations from unlabeled images to synthesize labeled examples for one-shot segmentation.
- Biomedical images vary widely in anatomy, contrast, and texture, complicating one-shot segmentation.
- Acquisition differences across machines and institutions produce variation in resolution, noise, and tissue appearance.
- Hand-tuned rotations and nonlinear deformations are easy to implement but have limited ability to emulate real image variations.
- The method learns spatial and appearance transformation models from unlabeled images and samples them to augment one labeled example.
- Training a supervised segmenter on synthesized labeled examples outperforms single-atlas segmentation and hand-tuned augmentation approaches.
2. Related work
Prior work addresses medical segmentation through atlas alignment, preprocessing, supervised augmentation, and few-shot methods, but often handles anatomy and intensity variation separately. This paper instead samples learned spatial and appearance transformations to create training examples for MRI segmentation.
- 2.1. Medical image segmentation: Brain MRI segmentation is challenged by anatomical variation and intensity changes caused by subjects, scanners, protocols, and image quality.
- 2.1. Medical image segmentation: The proposed augmentation improves robustness by retaining natural MRI intensity variations rather than removing them through preprocessing.
- 2.1. Medical image segmentation: Atlas-based methods align labeled reference volumes to targets and propagate labels through the deformation.
- 2.1. Medical image segmentation: Semi-supervised methods commonly use collections of segmentations, whereas this method leverages collections of unlabeled images.
- 2.2. Spatial and appearance transform models: The framework learns unconstrained spatial and intensity transformations and samples them to synthesize training examples rather than directly performing registration or segmentation.
- 2.3. Few-shot segmentation of natural images: Few-shot segmentation methods mainly target natural images, while medical images present different visual challenges.
- Parameterized augmentation can improve test performance, but gains depend on the selected transformation functions and parameters.
- Compared with earlier learned augmentation for digit classification, this work adds an appearance model and targets MRI segmentation.
3. Method
The method learns separate spatial and appearance transformation models from unlabeled volumes, then applies sampled transformations to one labeled atlas to create labeled training examples. These examples augment supervised segmentation training.
- The approach improves one-shot biomedical segmentation by synthesizing realistic labeled training examples in a semi-supervised framework.
- Transformation models: Separate models capture spatial and intensity differences between a labeled atlas and unlabeled volumes.Spatial transformations model anatomical variation, while appearance transformations model intensity variation.
- Transformation models: The spatial model estimates smooth deformations between the atlas and each unlabeled volume using a VoxelMorph variant.The variant uses normalized cross-correlation and a displacement-field smoothness term.
- Transformation models: The appearance model predicts per-voxel intensity changes after registering an unlabeled volume to atlas space.Its loss combines image similarity with semantically aware smoothness regularization based on anatomical boundaries.
- Synthetic examples: Sampling transformations from unlabeled target volumes can combine one subject’s spatial variation with another’s intensity pattern.Applying the spatial transformation to both the atlas image and labels produces correctly labeled synthetic examples for supervised training.
- Segmentation training: The synthesized labeled examples are added to the labeled training set for a supervised segmentation network trained on random image slices.The segmentation network is based on a state-of-the-art architecture and uses early stopping on a validation set.
4. Experiments
The experiments evaluate one-shot segmentation of unnormalized brain MRI using one labeled atlas, unlabeled training scans, validation scans, and a held-out test set.
- The study targets one-shot segmentation of unnormalized brain MRI, a challenging and practical setting.The authors note that intensity normalization can work poorly for clinical-quality scans and scans with stroke or traumatic brain injury.
- Data: The dataset contains T1-weighted MRI scans compiled from eight databases and evaluated across 30 anatomical labels.Brains are resampled to 1 mm isotropic voxels, affinely aligned, cropped, and not intensity-corrected.
- Data: Training uses 101 scans, but only one selected atlas has segmentation labels used for training the transformation models.An additional 50 scans form the validation set and 100 scans form the held-out test set.
4.2. Segmentation baselines
The evaluation compares single-atlas segmentation, SAS-based augmentation, hand-tuned random augmentation, and fully supervised segmentation using Dice performance on 100 held-out scans.
- SAS registers the atlas to each test volume and warps atlas labels using the computed deformation field.
- SAS-aug adds 100 unannotated training brains labeled using SAS results to the segmenter training set.
- Rand-aug synthesizes smooth deformations and intensity variations, with amplitude and smoothness selected for validation-set segmentation performance.
- The supervised method uses ground-truth labels for all 101 training examples and serves as an upper bound.
- Table 1 reports mean Dice, standard deviation, and mean pairwise improvement over SAS on 30 labels across 100 test subjects.
4.3. Variants of our method
The method includes independent and coupled sampling variants, plus a combination with hand-tuned augmentation, differing in how spatial and appearance transformations are paired.
- Ours-indep samples spatial and appearance target images independently, yielding 10,000 possible synthetic labeled examples from 100 unlabeled targets.Training samples one random labeled example per iteration because of memory constraints.
- Ours-coupled samples spatial and appearance transformations from the same target image, producing 100 possible synthetic examples.It also synthesizes a random example in each training iteration.
- Ours-indep + rand-aug alternates learned synthetic examples with hand-tuned augmentation during segmenter training.The added augmentation is intended to introduce variance unseen in the unlabeled set.
4.4. Evaluation metrics
Segmentation accuracy is measured with Dice score against anatomical labels generated using FreeSurfer.
- Dice quantifies overlap between two anatomical regions, ranging from 0 for no overlap to 1 for perfect overlap.
- Predicted segmentation labels are evaluated relative to anatomical labels generated using FreeSurfer.
4.5. Results
The learned augmentation methods improve segmentation over the evaluated baselines and show more consistent gains across subjects and anatomical structures.
- Our methods outperform all baselines in mean Dice score across all 30 evaluation labels, significantly exceeding rand-aug and SAS-aug.The reported paired t-test values are p < 1e-15 for rand-aug and p < 1e-20 for SAS-aug.
- Figures 4 and 5 compare each method’s pairwise mean-Dice improvement with the SAS baseline across test subjects.
- Ours-indep + rand-aug is consistently better than every baseline on every test subject.
- Ours-indep is better than rand-aug on 95 of 100 test scans and always better than SAS-aug and SAS.
- Rand-aug improves Dice for large structures but is detrimental for smaller ones, whereas the learned methods improve consistently across structures.
- Independent spatial and appearance models synthesize varied brain appearances with realistic images and accurate labels.
5. Discussion
The method improves segmentation by applying the same learned spatial transformation to images and labels, avoiding mismatched synthetic examples. Its demonstrated scope is brain MRI, with richer transformation sampling identified as a future extension.
- Why our method outperforms: Using identical spatial transformations for images and labels keeps synthesized training examples correctly labeled despite registration errors.SAS-aug instead applies warped labels to the original image, allowing registration errors or noise to create mislabeled examples.
- Why our method outperforms: The method synthesizes image texture within hippocampus labels that is more consistent with ground-truth texture, producing more useful training examples.
- Extensions: The framework currently samples transformations from a discrete set, with interpolation or composition proposed to cover the transformation space more richly.
- Extensions: The approach was demonstrated on brain MRIs, although its lack of brain- or MRI-specific information makes extension to other anatomy or modalities feasible.
6. Conclusion
The paper presents learning-based augmentation that uses one labeled image and unlabeled examples to synthesize realistic labeled training data. Applied to one-shot brain MRI segmentation, the synthesized examples outperform existing methods and approach fully supervised performance.
- 6. Conclusion: The paper presents a learning-based data augmentation method demonstrated on one-shot medical image segmentation.
- 6. Conclusion: Starting from one labeled image and unlabeled examples, the method learns spatial and appearance transformations and applies sampled transformations to synthesize realistic labeled images.
- 6. Conclusion: The method outperforms existing one-shot segmentation methods on every test example and approaches the performance of a fully supervised model.
- 6. Conclusion: Learning independent spatial and appearance transformation models from unlabeled images enables diverse and realistic labeled-example synthesis.
- 6. Conclusion: The synthesized examples can train a segmentation model that outperforms existing methods in a one-shot scenario.