Joint Optimization of Class-Specific Training- and Test-Time Data Augmentation in Segmentation

This paper presents an effective and general data augmentation framework for medical image segmentation. We adopt a computationally efficient and data-efficient gradient-based meta-learning scheme to explicitly align the distribution of training and validation data which is used as a proxy for unseen test data. We improve the current data augmentation strategies with two core designs. First, we learn class-specific training-time data augmentation (TRA) effectively increasing the heterogeneity within the training subsets and tackling the class imbalance common in segmentation. Second, we jointly optimize TRA and test-time data augmentation (TEA), which are closely connected as both aim to align the training and test data distribution but were so far considered separately in previous works. We demonstrate the effectiveness of our method on four medical image segmentation tasks across different scenarios with two state-of-the-art segmentation models, DeepMedic and nnU-Net. Extensive experimentation shows that the proposed data augmentation framework can significantly and consistently improve the segmentation performance when compared to existing solutions. Code is publicly available1.

[1]  D. Rueckert,et al.  Causality-Inspired Single-Source Domain Generalization for Medical Image Segmentation , 2021, IEEE Transactions on Medical Imaging.

[2]  Yizhou Yu,et al.  nnFormer: Interleaved Transformer for Volumetric Segmentation , 2021, ArXiv.

[3]  Aaron Carass,et al.  Autoencoder based self-supervised test-time adaptation for medical image analysis , 2021, Medical Image Anal..

[4]  Sergio Escalera,et al.  Multi-Centre, Multi-Vendor and Multi-Disease Cardiac Segmentation: The M&Ms Challenge , 2021, IEEE Transactions on Medical Imaging.

[5]  Yang Liu,et al.  Semi-supervised Cardiac Image Segmentation via Label Propagation and Style Transfer , 2020, M&Ms and EMIDEC/STACOM@MICCAI.

[6]  Jun Ma,et al.  Histogram Matching Augmentation for Domain Adaptation with Application to Multi-Centre, Multi-Vendor and Multi-Disease Cardiac Image Segmentation , 2020, M&Ms and EMIDEC/STACOM@MICCAI.

[7]  Konstantinos Kamnitsas,et al.  Analyzing Overfitting Under Class Imbalance in Neural Networks for Image Segmentation , 2020, IEEE Transactions on Medical Imaging.

[8]  Jens Petersen,et al.  nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation , 2020, Nature Methods.

[9]  Guha Balakrishnan,et al.  When and Why Test-Time Augmentation Works , 2020, ArXiv.

[10]  Peter M. Full,et al.  Studying Robustness of Semantic Segmentation under Domain Shift in cardiac MRI , 2020, M&Ms and EMIDEC/STACOM@MICCAI.

[11]  Younghoon Kim,et al.  Learning Loss for Test-Time Augmentation , 2020, NeurIPS.

[12]  Zhanxing Zhu,et al.  Automatic Data Augmentation for 3D Medical Image Segmentation , 2020, MICCAI.

[13]  Krishna Chaitanya,et al.  Semi-supervised Task-driven Data Augmentation for Medical Image Segmentation , 2020, Medical Image Anal..

[14]  Pheng-Ann Heng,et al.  Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains , 2020, MICCAI.

[15]  Daniel Rueckert,et al.  Realistic Adversarial Data Augmentation for MR Image Segmentation , 2020, MICCAI.

[16]  Víctor M. Campello,et al.  Cardiac Segmentation on Late Gadolinium Enhancement MRI: A Benchmark Study from Multi-Sequence Cardiac MR Segmentation Challenge , 2020, Medical Image Anal..

[17]  E. Konukoglu,et al.  Test-Time Adaptable Neural Networks for Robust Medical Image Segmentation , 2020, Medical Image Anal..

[18]  Timothy M. Hospedales,et al.  DADA: Differentiable Automatic Data Augmentation , 2020, ECCV.

[19]  Daguang Xu,et al.  Searching Learning Strategy with Reinforcement Learning for 3D Medical Image Segmentation , 2019, MICCAI.

[20]  Quoc V. Le,et al.  Randaugment: Practical automated data augmentation with a reduced search space , 2019, 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW).

[21]  F. Hutter,et al.  Understanding and Robustifying Differentiable Architecture Search , 2019, ICLR.

[22]  Taesup Kim,et al.  Fast AutoAugment , 2019, NeurIPS.

[23]  Adrian V. Dalca,et al.  Data Augmentation Using Learned Transformations for One-Shot Medical Image Segmentation , 2019, 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR).

[24]  Nassir Navab,et al.  Data Augmentation with Manifold Exploring Geometric Transformations for Increased Performance and Robustness , 2019, ArXiv.

[25]  Anant Gupta,et al.  Generative Image Translation for Data Augmentation of Bone Lesion Pathology , 2018, MIDL.

[26]  Sébastien Ourselin,et al.  Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks , 2018, Neurocomputing.

[27]  Yiming Yang,et al.  DARTS: Differentiable Architecture Search , 2018, ICLR.

[28]  Paolo Frasconi,et al.  Bilevel Programming for Hyperparameter Optimization and Meta-Learning , 2018, ICML.

[29]  Quoc V. Le,et al.  AutoAugment: Learning Augmentation Policies from Data , 2018, ArXiv.

[30]  Lei Ai,et al.  A large, open source dataset of stroke anatomical brain images and manual lesion segmentations , 2017, Scientific Data.

[31]  Hongyi Zhang,et al.  mixup: Beyond Empirical Risk Minimization , 2017, ICLR.

[32]  Graham W. Taylor,et al.  Improved Regularization of Convolutional Neural Networks with Cutout , 2017, ArXiv.

[33]  Sergey Levine,et al.  Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks , 2017, ICML.

[34]  Quoc V. Le,et al.  Neural Architecture Search with Reinforcement Learning , 2016, ICLR.

[35]  Ben Poole,et al.  Categorical Reparameterization with Gumbel-Softmax , 2016, ICLR.

[36]  Thomas Brox,et al.  3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation , 2016, MICCAI.

[37]  Nikos Komodakis,et al.  Wide Residual Networks , 2016, BMVC.

[38]  Ameet Talwalkar,et al.  Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization , 2016, J. Mach. Learn. Res..

[39]  Konstantinos Kamnitsas,et al.  Efficient multi‐scale 3D CNN with fully connected CRF for accurate brain lesion segmentation , 2016, Medical Image Anal..

[40]  Fabian Pedregosa,et al.  Hyperparameter optimization with approximate gradient , 2016, ICML.

[41]  Guillaume Lemaitre,et al.  Computer-Aided Detection and diagnosis for prostate cancer based on mono and multi-parametric MRI: A review , 2015, Comput. Biol. Medicine.

[42]  Jonathon Shlens,et al.  Explaining and Harnessing Adversarial Examples , 2014, ICLR.

[43]  Florian Jung,et al.  Evaluation of prostate segmentation algorithms for MRI: The PROMISE12 challenge , 2014, Medical Image Anal..

[44]  Luca Maria Gambardella,et al.  Mitosis Detection in Breast Cancer Histology Images with Deep Neural Networks , 2013, MICCAI.

[45]  Geoffrey E. Hinton,et al.  ImageNet classification with deep convolutional neural networks , 2012, Commun. ACM.

[46]  P. Cochat,et al.  Et al , 2008, Archives de pediatrie : organe officiel de la Societe francaise de pediatrie.

[47]  Xiao-Li Meng,et al.  The Art of Data Augmentation , 2001 .

[48]  Yoshua Bengio,et al.  Gradient-Based Optimization of Hyperparameters , 2000, Neural Computation.

[49]  S. K. Zhou,et al.  Transforming medical imaging with Transformers? A comparative review of key properties, current progresses, and future perspectives , 2022, Medical Image Anal..

[50]  Alex Krizhevsky,et al.  Learning Multiple Layers of Features from Tiny Images , 2009 .