UncertaintyFuseNet: Robust uncertainty-aware hierarchical feature fusion model with Ensemble Monte Carlo Dropout for COVID-19 detection

Abstract

The COVID-19 (Coronavirus disease 2019) pandemic has become a major global threat to human health and well-being. Thus, the development of computer-aided detection (CAD) systems that are capable of accurately distinguishing COVID-19 from other diseases using chest computed tomography (CT) and X-ray data is of immediate priority. Such automatic systems are usually based on traditional machine learning or deep learning methods. Differently from most of the existing studies, which used either CT scan or X-ray images in COVID-19-case classification, we present a new, simple but efficient deep learning feature fusion model, called $UncertaintyFuseNet$ , which is able to classify accurately large datasets of both of these types of images. We argue that the uncertainty of the model's predictions should be taken into account in the learning process, even though most of the existing studies have overlooked it. We quantify the prediction uncertainty in our feature fusion model using effective Ensemble Monte Carlo Dropout (EMCD) technique. A comprehensive simulation study has been conducted to compare the results of our new model to the existing approaches, evaluating the performance of competing models in terms of Precision, Recall, F-Measure, Accuracy and ROC curves. The obtained results prove the efficiency of our model which provided the prediction accuracy of 99.08% and 96.35% for the considered CT scan and X-ray datasets, respectively. Moreover, our $UncertaintyFuseNet$ model was generally robust to noise and performed well with previously unseen data. The source code of our implementation is freely available at: https://github.com/moloud1987/UncertaintyFuseNet-for-COVID-19-Classification.

Keywords: COVID-19; Deep learning; Early fusion; Feature fusion; Uncertainty quantification.

Fig. 1

A general overview of the applied deep learning model deep 1 (simple CNN).

Fig. 2

A general overview of the applied deep learning mode deep 2 (multi-headed CNN).

Fig. 3

A general overview of the proposed $UncertaintyFuseNet$ model inspired by a hierarchical feature fusion approach and EMCD.

Fig. 4

Some random image samples from the CT scan and X-ray datasets considered in our study.

Fig. 5

ROC curves obtained for the five considered ML models for the CT scan data without quantifying uncertainty.

Fig. 6

ROC curves obtained for the five considered ML models for the X-ray data without quantifying uncertainty.

Fig. 7

ROC curves obtained for the three considered DL models for the CT scan data with UQ.

Fig. 8

ROC curves obtained for the three considered DL models for the X-ray data with UQ.

Fig. 9

The MNIST sample image fed to the deep learning models as an unknown sample.

Fig. 10

T-SNE visualization of different models applied to the CT scan data without and with quantifying uncertainty.

Fig. 11

T-SNE visualization of different models applied to the X-ray data without and with quantifying uncertainty.

Fig. 12

Grad-CAM visualization for our proposed fusion model without and with UQ for nCT (Fig. 12, Fig. 12), NiCT (Fig. 12, Fig. 12), and pCT (Fig. 12, Fig. 12) classes using CT scan dataset.

Fig. 13

Grad-CAM visualization for our proposed fusion model without and with UQ for COVID-19 (Fig. 13, Fig. 13), Normal (Fig. 13, Fig. 13), and Pneumonia (Fig. 13, Fig. 13) classes using the X-ray dataset.

Fig. 14

The output posterior distributions of our proposed feature fusion model calculated for the nCT 14(a), NiCT 14(b) and pCT 14(c) data classes for the CT scan dataset, and the COVID-19 14(d), Normal 14(e) and Pneumonia 14(f) data classes for the X-ray dataset.

Fig. B.15

Confusion matrices obtained using different models for the CT scan datasets without quantifying uncertainty.

Fig. B.16

Confusion matrices obtained using different models for the X-ray dataset without quantifying uncertainty.

Fig. B.17

Confusion matrices obtained using different models for the CT scan dataset with quantifying uncertainty.

Fig. B.18

Confusion matrices obtained using different models for the X-ray dataset with quantifying uncertainty.