Source-linked AI summary
Automatic Brain Tumor Segmentation using Convolutional Neural Networks with Test-Time Augmentation
Guotai Wang, Wenqi Li, Sebastien Ourselin, Tom Vercauteren
TL;DR
Brain tumor segmentation is challenging, and limited medical imaging data motivates augmentation to improve CNN performance. This paper evaluates test-time augmentation across CNN architectures using spatial transformations and noise, finding improved segmentation accuracy and uncertainty estimation.
Problem
Limited medical imaging datasets make strong CNN performance difficult, while brain tumors have ambiguous boundaries and substantial variation in size, shape, and location.
Method
The paper applies test-time augmentation with multiple transformed inputs and shared network weights across 3D U-Net, multi-class WNet, and cascaded CNNs.
Results
3D UNet with test-time augmentation achieved Dice scores of 75.43%, 87.31% and 78.32% for enhancing tumor core, whole tumor and tumor core, respectively.
Takeaways & Limitations
Test-time augmentation improves segmentation accuracy across different CNN structures and provides uncertainty estimation for segmentation outputs.
Takeaways & Limitations
The study uses only flipping, rotation, and scaling for test-time augmentation because more complex transformations such as elastic deformations are less efficient for testing.
Abstract
from arXiv · showhide
Automatic brain tumor segmentation plays an important role for diagnosis, surgical planning and treatment assessment of brain tumors. Deep convolutional neural networks (CNNs) have been widely used for this task. Due to the relatively small data set for training, data augmentation at training time has been commonly used for better performance of CNNs. Recent works also demonstrated the usefulness of using augmentation at test time, in addition to training time, for achieving more robust predictions. We investigate how test-time augmentation can improve CNNs' performance for brain tumor segmentation. We used different underpinning network structures and augmented the image by 3D rotation, flipping, scaling and adding random noise at both training and test time. Experiments with BraTS 2018 training and validation set show that test-time augmentation helps to improve the brain tumor segmentation accuracy and obtain uncertainty estimation of the segmentation results.
1 Introduction
Brain tumor segmentation can support diagnosis, surgical planning, and treatment assessment, but ambiguous boundaries and large anatomical variation make automation difficult. CNNs and data augmentation address this challenge, and test-time augmentation is investigated for more robust predictions.
- Automatic segmentation can provide accurate and reproducible tumor delineation for diagnosis, surgical planning, and treatment assessment.
- Ambiguous tumor boundaries arise from smooth intensity gradients, partial volume effects, and bias field artifacts.
- Tumors vary substantially across patients in size, shape, and localization, limiting the usefulness of strong anatomical priors.
- Deep CNNs have achieved state-of-the-art performance for multi-modal brain tumor segmentation while learning features automatically from annotated images.
- Limited medical-imaging datasets motivate training-time augmentation, while combining predictions from transformed test images can improve robustness.
- The paper applies test-time augmentation to multi-class brain tumor segmentation across CNN structures to improve accuracy and estimate output uncertainty.
2 Methods
The method compares several CNN architectures and applies shared stochastic spatial and intensity augmentations during training and testing. Multiple test-time predictions are combined for segmentation, while their diversity supports pixel-wise uncertainty estimation.
- Network Structures: The study compares 3D U-Net, cascaded WNet/TNet/ENet networks, and a one-pass multi-class WNet.
- Network Structures: 3D U-Net uses downsampling and upsampling paths with shortcut connections and produces four segmentation-label channels.
- Network Structures: WNet uses anisotropic convolutions, dilated convolutions, residual connections, multi-scale prediction, and multi-view fusion to balance context, complexity, and memory.
- Data Augmentation for Training and Testing: Augmentation samples flipping, 3D rotation, scaling, and Gaussian intensity noise from specified prior distributions at both training and test time.
- Data Augmentation for Training and Testing: Test-time augmentation generates N transformed inputs through Monte Carlo sampling and combines their predictions using majority voting.
- Uncertainty Estimation: Pixel-wise uncertainty is estimated from prediction diversity using entropy rather than variance for the multi-modal segmentation output distribution.
3 Experiments and Results
Experiments on BraTS 2018 evaluated test-time augmentation (TTA) with multiple CNN architectures, showing more spatially consistent qualitative outputs and improved segmentation measurements across tumor structures.
- Experimental setup: BraTS 2018 experiments used images from 285 training patients and evaluated validation and testing cohorts with four MRI sequences.The sequences were T1, T1ce, T2 and FLAIR.
- Experimental setup: The study compared 3D UNet, multi-class WNet, and cascaded networks, with predictions generated using test-time augmentation.The cascaded approach separately segments whole tumor, tumor core and enhancing tumor core.
- Qualitative results: TTA made 3D UNet outputs more spatially consistent, smoothed multi-class WNet segmentations, and visually improved cascaded-network tumor-core results.The qualitative comparisons are shown for validation-set examples.
- Qualitative results: TTA filled an apparent 3D UNet tumor-core hole, reduced multi-class WNet over-segmentation, and improved spatial consistency in another example.The cascaded-network result also improved in this example.
- Quantitative results: 3D UNet Dice scores increased from 73.44%, 86.38% and 76.58% to 75.43%, 87.31% and 78.32% for enhancing tumor core, whole tumor and tumor core, respectively.These values are reported for the BraTS 2018 validation set.
- Quantitative results: Multi-class WNet Dice scores with TTA were 77.70%, 89.56% and 73.04% for enhancing tumor core, whole tumor and tumor core, respectively.The corresponding baseline scores were 75.70%, 88.98% and 72.53%.
- Quantitative results: On the BraTS 2018 testing set, cascaded networks with TTA achieved average Dice scores of 74.66%, 87.78% and 79.64% for enhancing tumor core, whole tumor and tumor core.The corresponding Hausdorff distances were 4.16mm, 5.97mm and 6.71mm.
4 Discussion and Conclusion
The discussion concludes that test-time augmentation improves segmentation accuracy across several CNN structures and provides uncertainty estimates, while identifying transformation efficiency and uncertainty modeling as boundaries for future work.
- Discussion: TTA using flipping, rotation and scaling improved segmentation accuracy for 3D UNet, multi-class WNet and cascaded networks.More complex elastic deformations were identified as possible but slower and less efficient for testing.
- Discussion: The method’s uncertainty estimates can support downstream tasks such as uncertainty-aware volume measurement and guiding user interactions.These applications are presented as potential uses of the uncertainty output.
- Future work: The study estimates test-image-based, or aleatoric, uncertainty but leaves epistemic uncertainty from CNN ensembles and combined methods for future work.The discussion suggests combining ensembles or test-time dropout with TTA for more comprehensive uncertainty analysis.
- Conclusion: The conclusion reports improved brain tumor segmentation accuracy and uncertainty estimation across different CNN structures using augmentation during training and testing.The tested augmentations included 3D rotation, flipping, scaling and random noise.