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 available.Comment: Accepted by IEEE Transactions on Medical Imagin