Source-linked AI summary
Semi-Supervised and Task-Driven Data Augmentation
Krishna Chaitanya, Neerav Karani, Christian Baumgartner, Olivio Donati, Anton Becker, Ender Konukoglu
TL;DR
Medical image segmentation requires large labelled datasets, but obtaining expert annotations is costly and small datasets can limit generalization. The paper learns task-driven deformation and intensity generators from labelled and unlabelled images to create corresponding augmented image-label pairs. On cardiac MRI segmentation with small training sets, the method substantially outperforms conventional augmentation techniques.
Problem
Small labelled datasets limit medical image segmentation, while obtaining large expert-annotated datasets is costly and time-consuming.
Method
The method learns conditional deformation and intensity transformations from labelled data, uses unlabelled images for distributional regularization, and optimizes the generators for segmentation.
Results
For cardiac MRI segmentation with small training datasets, the proposed augmentation method substantially outperforms conventional augmentation techniques.
Takeaways & Limitations
Task-driven augmentation can improve segmentation in small-labelled-data settings even when generated images are not visually hyper-realistic.
Abstract
from arXiv · showhide
Supervised deep learning methods for segmentation require large amounts of labelled training data, without which they are prone to overfitting, not generalizing well to unseen images. In practice, obtaining a large number of annotations from clinical experts is expensive and time-consuming. One way to address scarcity of annotated examples is data augmentation using random spatial and intensity transformations. Recently, it has been proposed to use generative models to synthesize realistic training examples, complementing the random augmentation. So far, these methods have yielded limited gains over the random augmentation. However, there is potential to improve the approach by (i) explicitly modeling deformation fields (non-affine spatial transformation) and intensity transformations and (ii) leveraging unlabelled data during the generative process. With this motivation, we propose a novel task-driven data augmentation method where to synthesize new training examples, a generative network explicitly models and applies deformation fields and additive intensity masks on existing labelled data, modeling shape and intensity variations, respectively. Crucially, the generative model is optimized to be conducive to the task, in this case segmentation, and constrained to match the distribution of images observed from labelled and unlabelled samples. Furthermore, explicit modeling of deformation fields allow synthesizing segmentation masks and images in exact correspondence by simply applying the generated transformation to an input image and the corresponding annotation. Our experiments on cardiac magnetic resonance images (MRI) showed that, for the task of segmentation in small training data scenarios, the proposed method substantially outperforms conventional augmentation techniques.
1 Introduction
The paper addresses medical image segmentation with small labelled datasets, where deep networks may fail to capture anatomical and imaging variations or generalize to shifted test distributions. It proposes semi-supervised, task-driven augmentation that uses unlabelled images and learned shape and intensity transformations.
- Motivation: Small annotated datasets make medical segmentation difficult because networks may miss anatomical or acquisition-related variations and generalize poorly to unseen images.Relevant variation includes subject-specific anatomy, lesions, acquisition protocols, and scanner differences.
- Proposed approach: The proposed method learns deformation-field and intensity transformations from labelled data to synthesize image-label pairs reflecting labelled and unlabelled image distributions.The transformations are designed to map available labelled images toward the broader pool of images.
- Proposed approach: The method additionally encourages generated image-label pairs to benefit the segmentation task rather than merely appear realistic.
- Contribution: The experiments report substantial improvements over existing augmentation and semi-supervised learning techniques for cardiac MRI segmentation.
- Related work: Prior augmentation methods include affine or elastic transformations, generative models trained on labelled data, and linear interpolation of images and labels.The related work also covers semi-supervised, weakly supervised, and other regularization methods for small labelled datasets.
2 Methods
The method jointly uses task-driven segmentation objectives and semi-supervised adversarial regularization to learn conditional deformation and intensity generators. These generators transform labelled images while preserving exact image-label correspondence.
- Data augmentation setup: Generated image-label pairs are formed by transforming labelled training examples, with the segmentation network trained on the resulting augmented set.The general augmentation setup includes original and generated image-label pairs.
- Semi-Supervised and Task-Driven Data Augmentation: The proposed optimization makes generated pairs useful for segmentation while regularizing them with image-distribution information from labelled and unlabelled images.The conditional generators are trained separately for deformation fields and intensity transformations.
- Deformation field generator: The deformation generator outputs dense per-pixel fields that warp an image and its corresponding one-hot label together.
- Additive intensity field generator: The intensity generator adds an intensity mask to an input image while leaving its segmentation mask unchanged.
- Regularization and optimization: The generators produce warped or intensity-shifted images, while an additional magnitude term discourages identity mappings and encourages non-trivial transformations.The deformation and intensity outputs are represented as v ◦ XL and ΔI + XL, respectively.
- Optimization sequence: After learning the task-specific generators, they are fixed and the segmentation network is reinitialized and retrained with original and augmented labelled data.
3 Dataset and Network details
The experiments use cardiac cine-MRI data with expert annotations for three cardiac structures and compare models under small labelled-data conditions. The system combines a segmentation network with generator and discriminator networks.
- Dataset: The dataset contains cardiac cine-MRIs from 100 subjects across five groups, with expert annotations for the left ventricle, myocardium, and right ventricle.Annotations cover end-systole and end-diastole images.
- Pre-processing: The data are bias-corrected, intensity-normalized, resampled to 1.367x1.367mm in-plane resolution, and cropped or padded before training.
- Network details: The proposed system has three network types: a segmentation network S, a generator G, and a discriminator D.The deformation and intensity generators share the described generator architecture except at the final layer.
- Network details: The generator combines image and 100-dimensional noise inputs through separate subnetworks before producing the transformation output.
- Network details: The discriminator uses convolutional and fully connected layers, while the segmentation network uses a U-net-like encoder-decoder architecture.
- Training details: Training uses weighted cross-entropy with empirically selected class weights, λadv=1, λbig=10^-3, batch size 20, and 10000 iterations.
4 Experiments
Experiments evaluate augmentation strategies for cardiac MRI segmentation with either one or three labelled training volumes, using Dice scores over three foreground structures. The comparisons include conventional, learned, Mixup, and adversarial-training baselines.
- Experimental setup: The experiments use fixed test, validation, and unlabelled sets while comparing settings with NL=1 and NL=3 labelled 3D training images.
- Baselines: Baselines include no augmentation, affine transformations, random elastic deformations, and random contrast or brightness fluctuations.
- Proposed variants: Learned augmentation variants separately use deformation fields, intensity fields, or both types of transformations.
- Additional comparisons: Additional comparisons apply Mixup to labelled images or to combinations involving original, affine-transformed, and learned-transformation images.
- Additional comparisons: The study also compares supervised and semi-supervised adversarial training methods on the same dataset.
- Evaluation: Performance is evaluated with the Dice similarity coefficient on 20 test subjects for the left ventricle, myocardium, and right ventricle.
5 Results and Discussion
The proposed augmentation substantially improves segmentation over other augmentation and semi-supervised methods, particularly with only one labelled 3D volume. Learned deformation and intensity transformations benefit from population-level variation and task-aware regularization.
- 5 Results and Discussion: The proposed method substantially improves performance over other data augmentation and semi-supervised adversarial learning methods, especially when only one 3D volume is used for training.The comparison is based on mean Dice scores over 15 runs, with visual support from Fig. 2.
- 5 Results and Discussion: Learned deformation-field augmentations outperform random elastic augmentations by spanning geometric variation across labelled and unlabelled images while remaining useful for segmentation.
- 5 Results and Discussion: Learned additive intensity masks improve over random intensity fluctuations by covering population intensity statistics while benefiting segmentation.
- 5 Results and Discussion: Combining learned deformation and intensity augmentations yields a substantial Dice-score improvement over using either augmentation alone.
- 5 Results and Discussion: Removing generator regularization preserves gains over random transformations but reduces their benefit, supporting the use of unlabelled images during augmentation learning.
6 Conclusion
The paper presents a semi-supervised, task-driven augmentation approach for cardiac MRI segmentation with as few as 1–3 labelled 3D images. It substantially outperforms conventional augmentation, even though generated images need not be visually hyper-realistic.
- 6 Conclusion: The approach learns conditional generative mappings between labelled and unlabelled images while mapping the corresponding segmentation annotations.
- 6 Conclusion: Task-dependent losses guide the generative models toward augmentations useful for segmentation rather than requiring visually hyper-realistic generated images.
- 6 Conclusion: The method substantially outperforms conventional data augmentation for cardiac MRI segmentation in settings with few labelled images.The conclusion targets training datasets containing as few as 1 to 3 labelled 3D images.