Elsevier

Neurocomputing

Volume 338, 21 April 2019, Pages 34-45
Neurocomputing

Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks

https://doi.org/10.1016/j.neucom.2019.01.103Get rights and content
Under a Creative Commons license
open access

Highlights

  • Different types of uncertainties for deep-learning based medical image segmentation were analysed.

  • We propose a general aleatoric uncertainty estimation method based on test-time augmentation.

  • A theoretical formulation of test-time augmentation was proposed.

  • The proposed method was validated with 2D fetal brain segmentation and 3D brain tumor segmentation tasks.

Abstract

Despite the state-of-the-art performance for medical image segmentation, deep convolutional neural networks (CNNs) have rarely provided uncertainty estimations regarding their segmentation outputs, e.g., model (epistemic) and image-based (aleatoric) uncertainties. In this work, we analyze these different types of uncertainties for CNN-based 2D and 3D medical image segmentation tasks at both pixel level and structure level. We additionally propose a test-time augmentation-based aleatoric uncertainty to analyze the effect of different transformations of the input image on the segmentation output. Test-time augmentation has been previously used to improve segmentation accuracy, yet not been formulated in a consistent mathematical framework. Hence, we also propose a theoretical formulation of test-time augmentation, where a distribution of the prediction is estimated by Monte Carlo simulation with prior distributions of parameters in an image acquisition model that involves image transformations and noise. We compare and combine our proposed aleatoric uncertainty with model uncertainty. Experiments with segmentation of fetal brains and brain tumors from 2D and 3D Magnetic Resonance Images (MRI) showed that 1) the test-time augmentation-based aleatoric uncertainty provides a better uncertainty estimation than calculating the test-time dropout-based model uncertainty alone and helps to reduce overconfident incorrect predictions, and 2) our test-time augmentation outperforms a single-prediction baseline and dropout-based multiple predictions.

Keywords

Uncertainty estimation
Convolutional neural networks
Medical image segmentation
Data augmentation

Cited by (0)

Guotai Wang obtained his Bachelor and Master degree of Biomedical Engineering in Shanghai Jiao Tong University in 2011 and 2014 respectively. He then obtained his PhD degree of Medical and Biomedical Imaging in University College London in 2018. His research interests include image segmentation, computer vision and deep learning.

Wenqi Li is a Research Associate in the Guided Instrumentation for Fetal Therapy and Surgery (GIFT-Surg) project. His main research interests are in anatomy detection and segmentation for presurgical evaluation and surgical planning. He obtained a BSc degree in Computer Science from the University of Science and Technology Beijing in 2010, and then an MSc degree in Computing with Vision and Imaging from the University of Dundee in 2011. In 2015, he completed his PhD in the Computer Vision and Image Processing group at the University of Dundee.

Michael Aertsen is a Consultant Pediatric Radiologist at University Hospitals of Leuven. He studied medecine at the University of Hasselt and the Katholieke Universiteit Leuven. He is specialized in fetal MRI and his main research focus is the fetal brain development with advanced MRI techniques.

Jan Deprest is a Professor of Obstetrics and Gynaecology at the Katholieke Universiteit Leuven and Consultant Obstetrician Gynaecologist at the University Hospitals Leuven (Belgium). He is currently the academic chair of his department and the director of the Centre for Surgical Technologies at the Faculty of Medicine. He established the Eurofoetus consortium, which is dedicated to the development of instruments and techniques for minimally invasive fetal and placental surgery.

Sébastien Ourselin is Head of the School of Biomedical Engineering & Imaging Sciences and Professor of Healthcare Engineering at Kings College London. His core skills are in medical image analysis, software engineering, and translational medicine. He is best known for his work on image registration and segmentation, its exploitation for robust image-based biomarkers in neurological conditions, as well as for his development of image-guided surgery systems.

Tom Vercauteren is a Professor of Interventional Image Computing at Kings College London. He is a graduate from Columbia University and Ecole Polytechnique and obtained his PhD from Inria Sophia Antipolis. His main research focus is on the development of innovative interventional imaging systems and their translation to the clinic. One key driving force of his work is the exploitation of image computing and the knowledge of the physics of acquisition to move beyond the initial limitations of the medical imaging devices that are developed or used in the course of his research.