Bayesian deep learning for uncertainty aware medical image segmentation correction and feature engineering: Applications to OCT imaging

The clinical integration of AI in medical imaging faces challenges in generalizability, robustness, and trust, largely due to a lack of uncertainty quantification, resulting in “confidently incorrect” outputs. This paper introduces a novel framework using Bayesian deep learning, specifically Bayes by Backprop, for optical coherence tomography (OCT) image segmentation including uncertainty-aware corrections and feature engineering. Our approach provides precise segmentations along with epistemic and aleatoric uncertainty estimates; these uncertainties are valuable for identifying out-of-distribution (OOD) samples and facilitating automated segmentation corrections, thereby enhancing AI model reliability and interpretability in clinical contexts. Comparative analyses show our Bayesian models achieve comparable or marginally superior performance to deterministic counterparts, while offering significant qualitative gains in transparency and explainability, fostering greater clinician confidence. This work also demonstrates uncertainty-aware feature engineering by leveraging probabilistic Bayesian neural network outputs to derive robust clinical features, such as retinal layer volumes, complete with associated uncertainty estimates. This enables more sophisticated downstream statistical analyses, including mixed-effects modelling which can differentiate between model-induced variance and inter-sample variability. Ultimately, this research contributes to bridging the gap between advanced AI methodologies and their practical, trustworthy application in medical imaging.