Bayesian Uncertainty Estimation Improves nnU-Net Generalization to Unseen Sites for Stroke Lesion Segmentation

Abstract

This work studies the effectiveness of Bayesian uncertainty estimation in neural networks for enhancing domain generalization in stroke lesion segmentation. We employ a multi-modal posterior sampling approach for nnU-Net and evaluate it against the conventional nnU-Net. Both models are trained on data from one clinical site and tested on another for the most common imaging in stroke: NCCT, CTA, and CTP. The models produce segmentation probabilities for stroke core and hypoperfused tissue, which are thresholded at different levels and compared to ground truth labels to measure correctly predicted, missing and excess volume, and Dice scores. Conventional nnU-Net is limited in sensitivity adjustment, while the Bayesian approach allows for a flexible trade-off between under- and over-segmentation. The Bayesian nnU-Net demonstrates similar or superior performance w.r.t. Dice scores, with the added advantage of providing more insight into uncertain areas, enabling more nuanced decision-making. The Bayesian approach proved robust in handling data from unseen domains, suggesting its potential for enhancing automated stroke lesion detection across various clinical environments.