Uncertainty-Guided Self-Learning Framework for Semi-Supervised Multi-Organ Segmentation

Natália Alves, Bram de Wilde

Keywords: Semi-supervised learning, Uncertainty, Pseudo-labeling

TL;DR: Incorporating prediction uncertainty to generate high qualitiy pseudo-labels for semi-supervised medical image segmentation.

Abstract:

Automatic multi-organ segmentation in medical imaging has important clinical applications, but manual voxel-level annotations are time and labour-consuming, limiting the annotated data available for training. We propose an uncertainty-guided framework for multi-organ segmentation on CT scans that uses a small labelled dataset to leverage a large unlabelled dataset in a semi-supervised setting. First, we train five models to segment 13 abdominal organs using 50 manually labelled training cases and 5-fold cross-validation. Then, we use these models to generate pseudo-labels for 2000 unlabelled cases and estimate the uncertainty associated with the pseudo-labels by calculating the pair-wise Dice score (DSC) for the five individual predictions. Cases with pair-wise mean DSC>0.9 for all organs are included in the training set at the next iteration, together with the respective pseudo-labels.
This process is repeated for four iterations. All selected cases are combined in the last iteration, and a single model is trained to reduce the computational costs associated with ensembling. The self-configuring method for biomedical image segmentation nnU-Net was used to train the segmentation models. We obtained a mean DSC of 0.8388 on the validation set with the network trained using the labelled data alone. The Dice score improved to 0.8874 in the final iteration of the model, trained with the 50 labelled cases and 1813 unlabelled cases with pseudo-labels. On the final test set, the mean DSC was 0.8685, and the mean inference time per case was 42 seconds. All code is open-source and available on \href{https://github.com/DIAGNijmegen/flare22-brananas}{GitHub} (https://github.com/DIAGNijmegen/flare22-brananas).