Exploiting probability density function of deep convolutional autoencoders’ latent space for reliable COVID-19 detection on CT scans

There is general agreement on chest computed tomography (CT) scans having a potential role in diagnosing COVID-19 and being particularly effective when used as a complement to polymerase chain reaction (PCR) testing [1‐5]. Clearly, cheaper testing methods exist, like lateral flow test. Moreover, the radiation dose inflicted to the patient by a scan may be harmful. However, the method can be useful when a different test is unavailable or too expensive or when a quick diagnosis is needed.

Due to its noteworthy detection capabilities, deep learning (DL) is often employed to assist the evaluation of CT scans for the diagnosis of COVID-19. Indeed, classification of chest CT scans is an important and active research area. In recent contributions, classification is carried out by exploiting supervised DL, using only a few classes of chest CTs (for instance, COVID-19 class and normal one). However, the symptoms of chest diseases manifest themselves in a broad spectrum of visual characteristics. Hence, supervised-trained models could run into trouble when they are fed with chest CTs that do not belong to any of the classes used in the training phase. This issue can be addressed by taking advantage of unsupervised DL, where the neural network models are trained only on data sets of the COVID-19 class. In this manner, the trained model can differentiate the COVID-19 class (the target class) from any other types of chest images (anomalies).

This paper develops an unsupervised classification approach based on autoencoders (AEs). An AE is composed of an encoder and a decoder [6]. The encoder takes an image and transforms it into a hidden feature vector, a compressed input representation. The decoder is used to reconstruct the input image from the hidden feature vector. Since deep convolutional neural networks (CNNs) have proved their effectiveness in image processing and feature extraction, in this paper, we design a deep convolutional autoencoder (DCAE) as the neural network model [7].

We train the proposed DCAE, in an unsupervised manner, on a data set of chest CT scans obtained from COVID-19 patients. After training, the encoder is used to obtain the hidden feature vectors of the training set CT scans. These hidden feature vectors are, in turn, exploited to estimate the probability density function (PDF) of COVID-19 hidden feature vectors, by means of the multivariate kernel density estimation (KDE) method. The classification of a test image is performed by feeding it into the trained encoder to produce its hidden feature vector. Afterwards, the PDF is used to compute the PDF value of the test image. Finally, the obtained PDF value is compared to a suitably tuned threshold to classify the test image as either COVID-19 or non-COVID-19.

For comparison, we also consider a second method for classification based on the DCAE reconstruction error. The resulting error is noticeably higher than the average error corresponding to the training set instances when the trained DCAE is fed with a CT scan of a non-COVID-19 case and attempts to recover it. If this error is below a suitably tuned threshold, the image is classified as COVID-19, otherwise as non-COVID-19.

The rest of the paper is organised as follows: In Sect. 2, we review the related work and present the contributions of this paper. In Sect. 3, we provide details of the used data sets and pre-processing tasks. Sections 4 and 5 describe the proposed DCAE architecture and classification approach, respectively. We discuss the obtained numerical results in Sect. 6 and compare the performance of our approach to the ones of some state-of-the-art approaches that rely on the supervised training method. Finally, in Sect. 7, we point out the summarised observations and provide some hints for future research.

The literature on the application of DL-based algorithms to the detection of COVID-19 is vast. The small volume of available data on COVID-19 patients has motivated the researchers to take this deficiency into account. For instance, the transfer learning approach is adopted in [8‐17] to deal with the lack of large-size data sets. In [18], the authors utilise GoogleNet and ResNet for supervised COVID-19 classification. The authors of [19] take a statistical method to address issues like huge computational complexity and large datasets required by deep networks. The adopted approach is based on the evaluating and comparing the statistical representation of medical images. The authors of [20] consider an unbalanced-data supervised algorithm and obtained good results comparable with benchmark architectures.

A large number of research papers adopt supervised learning approaches [21‐31]. In [21], the authors consider a binary classification problem and apply the off-the-shelf VGG-16. In [26], the authors use depth-wise convolutions with varying dilation rates to extract more diversified features of chest images. The authors use a pre-trained model and reach the overall \(90.2 \%\) accuracy. In [28], the authors design a neural network model as a combination of convolutional and capsule layers, called COVID-FACT. Despite their great effort, the considered model achieves the \(90.82 \%\) accuracy. The authors of [29] propose a model based on the pre-trained ResNet50 and achieve the accuracy, rather similar to the original ResNet50. The DenseNet-based approach, considered in [31], achieves \(92\%\) accuracy. The work by [32] compares different classification architectures on a specific dataset to discover the most suitable real-world scenarios.

The approaches developed in all these papers typically deal with only a limited class of chest images. Hence, a naturally posed question is: to which category does a CT fit when it does not belong to any of the classes learned by the supervised-trained models? In any case, the aforementioned supervised methods must be trained on both COVID-19 and non-COVID-19 CTs.

Lastly, we provide an overview of AE-based approaches. The authors of [33] consider a stacked AE (SAE), composed of four CAEs, followed by a dense layer, and a final softmax classifier. Each layer is equipped with regularisation to improve the local optimum. The binary classification task, considered by the authors, occurs at the last stage, where the softmax classifier obtains the probability of the two types of labels and performs the classification task with an average accuracy of \(94.7 \%\). This method is different from our approach because we perform the classification by directly using the PDF of the hidden feature vectors, without inserting any classifier at the top of our model. The deep convolutional denoising AE, proposed in [34], is trained on COVID-19, pneumonia, and a few other types of chest X-rays. Then, the hidden feature vector of a test image is compared to the features of the selected training data sets. The considered AE exhibits good performance. However, unlike our work, this approach relies on training the considered model over each selected class and therefore cannot detect chest CTs except those that belong to the classes of training data sets. The work by [35] focuses on a two-stage learning method and a triple classification task. The authors train their considered model on classes of COVID-19, pneumonia, and normal cases separately. Once the hidden feature vectors of all classes are independently obtained, a feature classifier is employed and trained—in a supervised manner—to detect each decision class. The considered approach reaches a quite good accuracy of \(93.50 \%\). In contrast to this work, we train our DCAE model on only one class, i.e. the COVID-19. The paper by [36] is based on a variational autoencoder (VAE) model for COVID-19 classification. The VAE model involved adaptive Wiener filtering (AWF)-based pre-processing technique to enhance the image quality. Besides, Inception v4 with Adagrad technique is employed as a feature extractor and unsupervised VAE model is applied for the classification process. As the last research paper, the method introduced by [37] builds a robust statistical generating target histogram of the deep denoising convolutional autoencoder’s (DDCAE) latent vector. It then estimates the statistical distance between unknown and target histograms to classify the images according to proper thresholds.

A brief overview of the given literature is provided in Table 1. Finally, the main contributions of this paper are listed below:

The training data set used in this paper is composed of 4000 CTs of COVID-19 cases collected from over 500 patients. These CT scans have been selected from the ‘COVIDx CT-2A’ data set; the ‘A’ variant of the ‘2nd’ version of the open-source ‘COVIDx-CT’ data set [38]. The ‘COVIDx CT-2A’ data has been validated by medically qualified personnel. The data are split into training and validation sets containing \(80\%\) and \(20\%\) of instances, respectively.

The test data set comprises 4000 CT slices of COVID-19, normal, pneumonia, and three types of lung cancers, namely adenocarcinoma, large cell carcinoma and squamous cell carcinoma. The test images have been randomly selected from two separate data sets [38, 39], composed of CTs of over 500 patients different from those involved in the training set. The data sets include both male and female patients and cover a wide age range.

The information about the employed training/test data sets is summarised in Table 2, where we also give the links to the corresponding web pages. For illustrative purpose, in Fig. 1, we present some instances of various chest CT categories, which are drawn from training and test data sets. As can be seen from the figure, the COVID-19 image is characterised by the presence of several opacities. Indeed, commonly reported imaging features specific of COVID-19 pneumonia are peripheral, bilateral, ground-glass opacities with or without visible intralobular lines: see [40] for a wider discussion.

As a pre-processing task, the margin of every CT is cropped before performing the training phase. Moreover, all the images are resized to 300 pixels in width and 200 pixels in height: this size was selected as a compromise between computational complexity and resolution. Two samples of the COVID-19 dataset in the original, cropped and resized versions are demonstrated in Fig. 2.

A DCAE is composed of an encoder and a decoder as depicted in Fig. 3. The encoder takes an image as input and transforms it into a hidden feature vector called latent space. The decoder takes the hidden feature vector and attempts to recover the input image. The difference between the input and the output images results in the reconstruction error, which is the cost function to minimise during the training phase. We point out that the hidden feature vectors have fewer dimensions than those of the input images since the DCAE produces compressed versions.

The DCAE, exploited in this work, has been designed to take account of the trade-off between test classification accuracy and model complexity. Its “ad hoc” designed architecture is detailed in Table 3: the encoder is composed of three convolutional layers, two batch normalisation layers, two max-pool layers, followed by flatten and dense layers with rectifier activation functions (ReLUs).

The DCAE is trained to recover its input, i.e. COVID-19 CT scans. As training cost function, we use the mean squared error (MSE), and in order to minimise the considered cost function over the training set, we employ the Adam (Adaptive moment estimation) solver, which is a gradient-based stochastic optimisation algorithm that takes account of the first and second moments of underlying gradients [41]. The training is carried out over 100 epochs. After the training phase, we select the DCAE weights, which give rise to the minimum squared error on the validation data set. Table 4 recap details on the adopted model setting. For illustrative purpose, Fig. 4 presents the graphs of MSE and accuracy versus the number of epochs, over which the training phase is carried out.

Note that we have considered other options for both the architecture and the training parameters, using several alternatives proposed in the literature. The various options have been compared by means of unreported tests: the final architecture and training parameters have been selected because they yielded the highest classification accuracy.

The DCAE is implemented in Python language, using TensorFlow and Keras API. All the numerical trials have been carried out on a PC equipped with an AMD Ryzen 9 5900X 12-Core 3.7 GHz processor, two GeForce RTX 3070 graphics cards, and 128 GB RAM.

In this section, we describe the classification approach that is based on the PDF estimation of the DCAE hidden feature vectors.

In order to evaluate the performance of our classification method, we carried out several tests. We split the presentation into three parts that are: (i) the results that are obtained from the KDE probabilistic approach; (ii) the results that are achieved through reconstruction error evaluation; and (iii) performance comparisons with some state-of-the-art solutions. In the next subsection, we describe the employed performance metrics as a preliminary step.

We propose a method for classifying lung CTs as COVID-19 or non-COVID-19. The method exploits a DCAE trained on COVID-19 CTs only and a KDE estimation of the PDF of the DCAE hidden feature vectors. The DCAE is used to produce the corresponding hidden feature vector to classify an image. Afterwards, we use the so-obtained PDF evaluation to compute, in the test phase, the PDF value of the hidden feature vector, corresponding to a test image: if the PDF value is above a suitable threshold, that image is classified as COVID-19, otherwise as non-COVID-19.

We compare our KDE-based approach to the benchmark method that is based on the reconstruction error. In addition, we also check the accuracy performance of three widely known supervised models and the results of some recent papers. The carried out tests support the conclusion that the proposed approach is highly effective in terms of both achieved test accuracy and the needed test times.

The presented results could be extended, at least, along two main research lines. A first research line could concern the utilisation of generative adversarial networks (GANs) for the generation of additional training examples in the case of new COVID-19 mutations (as, for example, the Omicron one), in order to quickly provide reliable automatic detection of these mutations without pausing for the acquisition of sufficiently large new datasets. A second research line could regard the implementation of the proposed algorithmic framework atop distributed Fog/Cloud networked technological platforms [50, 51], in order to be capable to quickly generate reliable clinical diagnosis by leveraging the low-delay and (possibly, adaptive [52] and /or multiantenna empowered [53, 54]) capability of emerging Fog computing platforms for allowing ubiquitous wireless access to computing-demanding medical diagnostic services.