Active Learning with the nnUNet and Sample Selection with Uncertainty-Aware Submodular Mutual Information Measure

30 Jan 2024 (modified: 21 Mar 2024)MIDL 2024 Conference SubmissionEveryoneRevisionsBibTeXCC BY 4.0
Keywords: Deep Learning, Active Learning, Medical Image Segmentation, nnUNet, Submodular Subset Selection
Abstract: Annotating medical images for segmentation tasks is a time-consuming process that requires expert knowledge. Active learning can reduce this annotation cost and achieve optimal model performance by selecting only the most informative samples for annotation. However, the eectiveness of active learning sample selection strategies depends on the model architecture and training procedure used. The nnUNet has achieved impressive results in various automated medical image segmentation tasks due to its self-configuring pipeline for automated model design and training. This raises the question of whether the nnUNet is applicable in an active learning setting to avoid cumbersome manual configuration of the training process and improve accessibility for non-experts in deep learning-based segmentation. This paper compares various sample selection strategies in an active learning setting in which the self-configuring nnUNet is used as the segmentation model. Additionally, we propose a new sample selection strategy for UNet-like architectures: USIM - Uncertainty-Aware Submodular Mutual Information Measure. The method combines uncertainty and submodular mutual information to select batches of uncertain, diverse, and representative samples. We evaluate the performance gain and labeling costs on three medical image segmentation tasks with different segmentation challenges. Our findings demonstrate that utilizing nnUNet as the segmentation model in an active learning setting is feasible, and most sampling strategies outperform random sampling. Furthermore, we demonstrate that our proposed method yields a significant improvement compared to existing baseline methods.
Submission Number: 146