A sensitivity analysis of probability maps in deep-learning-based anatomical segmentation

Abstract

Purpose: Deep-learning-based segmentation models implicitly learn to predict the presence of a structure based on its overall prominence in the training dataset. This phenomenon is observed and accounted for in deep-learning applications such as natural language processing but is often neglected in segmentation literature. The purpose of this work is to demonstrate the significance of class imbalance in deep-learning-based segmentation and recommend tuning of the neural network optimization objective. Methods: An architecture and training procedure were chosen to represent common models in anatomical segmentation. A family of 5-block 2D U-Nets were independently trained to segment 10 structures from the Cancer Imaging Archive's Head-Neck-Radiomics-HN1 dataset. We identify the optimal threshold for our models according to their Dice score on the validation datasets and consider perturbations about the optimum. A measure of structure prominence in segmentation datasets is defined, and its impact on the optimal threshold is analyzed. Finally, we consider the use of a 2D Dice objective in addition to binary cross entropy. Results: We observe significant decreases in perceived model performance with conventional 0.5-thresholding. Perturbations of as little as ±0.05 about the optimum threshold induce a median reduction in Dice score of 11.8% for our models. There is statistical evidence to suggest a weak correlation between training dataset prominence and optimal threshold (Pearson (Formula presented.) and (Formula presented.)). We find that network optimization with respect to the 2D Dice score itself significantly reduces variability due to thresholding but does not unequivocally create the best segmentation models when assessed with distance-based segmentation metrics. Conclusion: Our results suggest that those practicing deep-learning-based contouring should consider their postprocessing procedures as a potential avenue for improved performance. For intensity-based postprocessing, we recommend a mixed objective function consisting of the traditional binary cross entropy along with the 2D Dice score.