Uncertainty-aware self-training with adversarial data augmentation for semi-supervised medical image segmentation

Abstract

Supervised algorithms require a significant amount of labeled data to ensure the effectiveness and robustness. Unfortunately, obtaining segmentation masks annotated by experts is both time-consuming and expensive. Although existing methods use data augmentation to expand the training data, these approaches often only slightly improve the generalization. To address this issue, we proposes a medical image segmentation framework that aims to leverage unlabeled samples for feature learning and improve segmentation performance. The proposed framework includes a data augmentation model and a segmentation model. The data augmentation model utilizes generative adversarial networks to model the spatial and intensity transformations in medical images and generate strongly-augmented samples to expand the training set. The segmentation model is implemented by applying self-training methods and consistency regularization. Firstly, pseudo-labeling is performed on weakly-augmented samples. Then, consistency regularization is applied to encourage the model predictions on strongly-augmented samples to be consistent with the pseudo-labels. This aims to improve the robustness of the model on unseen samples. To mitigate the network degradation caused by unreliable pseudo-labels, a new self-training strategy and uncertainty estimation are introduced into the segmentation framework to enhance the reliability of pseudo-labels. The proposed framework is rigorously evaluated for the segmentation of cardiac and prostate images, the experimental results indicate that it achieves competitive performance compared to several state-of-the-art methods. Moreover, the proposed method is applicable for joint training with limited labeled and additional unlabeled data, potentially reducing the workload of obtaining annotated images.