Uncertainty-Aware Deep Learning for 3D Fetal Brain Pose Prediction from Freehand 2D Ultrasound
Poster No:
93
Submission Type:
Abstract Submission
Authors:
Jayroop Ramesh1, Pak Hei Yeung2, Ana Namburete1
Institutions:
1University of Oxford, Oxford, Oxfordshire, 2University of Oxford, OXFORD, Oxfordshire
First Author:
Co-Author(s):
Introduction:
Two-dimensional (2D) ultrasound (US) is the preferred tool for routine monitoring and assessment of fetal growth and anatomy [1]. By capitalizing on the availability and flexibility of low-cost freehand 2D US scanners, it is possible to provide routine prenatal monitoring in resource-constrained settings. However, scanning is heavily operator-dependent, and scarcity of skilled sonographers limits its use. Deep Neural Networks (DNNs) can assist with automated US analysis [2], but their robustness is affected by variability in quality of US images collected by different operators with subjective acquisition protocols. We, therefore, propose an uncertainty-aware deep learning model for 3D pose prediction of 2D fetal brain images, to be used for scanning guidance. Specifically, we train a multi-head network to jointly regress 3D plane poses from 2D images in terms of different geometric transformations and their respective data-dependent uncertainties. Leveraging the output uncertainties can result in a model that is more robust to noise effects observed in freehand US scanning.
Methods:
Our data consisted of 3D US volumes acquired at 19 gestational weeks as part of the INTERGROWTH-21st study [3] and aligned to a reference 3D atlas space [4]. We selected 24, 2, and 7 3D volumes for training, validation, and testing respectively. The proposed network took as input 2D slices sampled from arbitrary cross-sectional planes of the 3D volumes. The corresponding plane poses are parameterized by xyz coordinates of 3 reference points and defined as 3D pose P∈Rheight×width×9.
We adapt existing work [2] which predicts the 3D pose P¬ of 2D US fetal brain images by incorporating components to account for uncertainty. We refer to our proposed model as QAERTS, which predicts the pose represented by various parameterizations of rotation (i.e., Quaternions, Axis-angles, Euler angles, Rotation matrices) in addition to shared Translation and Scaling using a multi-head DNN (Fig 1). We hypothesize that confidence of predictions with inputs of variable quality can be quantified by measuring variance between different parameterizations of predicted poses.
As the original loss function of MSE in [2] does not capture predictive uncertainty and assumes uniform variance across all inputs, we utilize the outputs to predict a multivariate normal distribution parameterized by the ensemble mean (P¬avg) of poses and the learned variances (σ2avg) after each geometrical transformation. Then, Gaussian Negative Likelihood Loss (GNLL) is minimized with respect to ground-truth reference poses. Accounting for heteroscedasticity during training allows higher weight to be assigned to inputs with lower variance, and improves learning by focusing on lower-noise regions in feature space [6].
The evaluation metrics used are Euclidean distance (ED), plane angle (PA), normalized cross-correlation (NCC) and structural similarity (SSIM) [2].
We adapt existing work [2] which predicts the 3D pose P¬ of 2D US fetal brain images by incorporating components to account for uncertainty. We refer to our proposed model as QAERTS, which predicts the pose represented by various parameterizations of rotation (i.e., Quaternions, Axis-angles, Euler angles, Rotation matrices) in addition to shared Translation and Scaling using a multi-head DNN (Fig 1). We hypothesize that confidence of predictions with inputs of variable quality can be quantified by measuring variance between different parameterizations of predicted poses.
As the original loss function of MSE in [2] does not capture predictive uncertainty and assumes uniform variance across all inputs, we utilize the outputs to predict a multivariate normal distribution parameterized by the ensemble mean (P¬avg) of poses and the learned variances (σ2avg) after each geometrical transformation. Then, Gaussian Negative Likelihood Loss (GNLL) is minimized with respect to ground-truth reference poses. Accounting for heteroscedasticity during training allows higher weight to be assigned to inputs with lower variance, and improves learning by focusing on lower-noise regions in feature space [6].
The evaluation metrics used are Euclidean distance (ED), plane angle (PA), normalized cross-correlation (NCC) and structural similarity (SSIM) [2].
Results:
As shown in Fig 1(b-e) and 2, our proposed model, QAERTS, with ablation studies, were compared to Base [2], and its modification to predict mean and variance (MVE) from separate heads [6], as well as current uncertainty-based deep learning baselines, namely Monte-Carlo Dropout (MCD) [7], Deep ensembles (DE) [8] and Deep Evidential Regression (EDL) [9].
QAERTS reports improved performances across each metric compared to base, MVE, MCD and EDL, but is not as capable as DE (Fig 1b-e) [10]. Nevertheless, QAERTS mitigates computational overhead in terms of time and parameters compared to DE, while maintaining competitive performance on pose prediction quality compared to other baselines (Fig 2).
QAERTS reports improved performances across each metric compared to base, MVE, MCD and EDL, but is not as capable as DE (Fig 1b-e) [10]. Nevertheless, QAERTS mitigates computational overhead in terms of time and parameters compared to DE, while maintaining competitive performance on pose prediction quality compared to other baselines (Fig 2).
Conclusions:
We observe overall performance of all techniques improves through uncertainty-aware learning and our proposed model, QAERTS, was second only in predictive performance to DE but with ~5x fewer parameters. This suggests QAERTS can reduce model ambiguity with respect to input quality while being computationally efficient.
Brain Stimulation:
Sonic/Ultrasound 1
Modeling and Analysis Methods:
Classification and Predictive Modeling 2
Keywords:
Data analysis
Machine Learning
ULTRASOUND
1|2Indicates the priority used for review
Provide references using author date format
[2] P.-H. Yeung, M. Aliasi, A. T. Papageorghiou, M. Haak, W. Xie, and A. I. L. Namburete, ‘Learning to map 2D ultrasound images into 3D space with minimal human annotation’, Med. Image Anal., vol. 70, p. 101998, May 2021, doi: 10.1016/j.media.2021.101998.
[3] A. T. Papageorghiou et al., ‘The INTERGROWTH-21st fetal growth standards: toward the global integration of pregnancy and pediatric care’, Am. J. Obstet. Gynecol., vol. 218, no. 2S, pp. S630–S640, Feb. 2018, doi: 10.1016/j.ajog.2018.01.011.
[4] A. I. L. Namburete et al., ‘Normative spatiotemporal fetal brain maturation with satisfactory development at 2 years’, Nature, pp. 1–9, Oct. 2023, doi: 10.1038/s41586-023-06630-3.
[5] A. Kendall and R. Cipolla, ‘Geometric Loss Functions for Camera Pose Regression with Deep Learning’. arXiv, May 23, 2017. doi: 10.48550/arXiv.1704.00390.
[6] D. A. Nix and A. S. Weigend, ‘Estimating the mean and variance of the target probability distribution’, in Proceedings of 1994 IEEE International Conference on Neural Networks (ICNN’94), Jun. 1994, pp. 55–60 vol.1. doi: 10.1109/ICNN.1994.374138.
[7] Y. Gal and Z. Ghahramani, ‘Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning’. arXiv, Oct. 04, 2016. Accessed: May 15, 2023. [Online]. Available: http://arxiv.org/abs/1506.02142
[8] B. Lakshminarayanan, A. Pritzel, and C. Blundell, ‘Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles’.arXiv, Dec. 05, 2016. doi: https://arxiv.org/abs/1612.01474.
[9] A. Amini, W. Schwarting, A. Soleimany, and D. Rus, ‘Deep Evidential Regression’. arXiv, Nov. 24, 2020. doi: 10.48550/arXiv.1910.02600.
[10] M. Havasi et al., ‘Training independent subnetworks for robust prediction’. arXiv, Aug. 04, 2021. doi: 10.48550/arXiv.2010.06610.