COVID-FACT: A Fully-Automated Capsule Network-based Framework for Identification of COVID-19 Cases from Chest CT scans

Convolutional Neural Networks (CNNs) have been widely used in several studies to account for the human-centered weaknesses in detecting COVID-19. CNNs are powerful models in related tasks and are capable of extracting distinguishing features from CT scans and chest radiographs (Yamashita et al., 2018). In this regard, many studies have utilized CNNs to identify COVID-19 cases from medical images. The study by (Wang and Wong, 2020), is an example of the application of CNN in COVID-19 detection, where CNN is first pre-trained on the ImageNet dataset (Krizhevsky et al., 2017). Fine-tuning is then performed using a CR dataset. Results show an accuracy of 93.3% in distinguishing normal, non-COVID-19 pneumonia (viral and bacterial), and COVID-19 infection cases. (Sethy et al., 2020) have also explored the same problem with the difference that the CNN is followed by a Support Vector Machine (SVM), to identify positive COVID-19 cases. Their obtained results show an overall accuracy of 95.38%, sensitivity of 97.29% and specificity of 93.47%. Another study by (Mahmud et al., 2020) proposed a CNN-based model utilizing depth-wise convolutions with varying dilation rates to extract more diversified features from chest radiographs. They used a pre-trained model on a dataset of normal, viral, and bacterial pneumonia patients followed by additional fine-tuned layers on a dataset of COVID-19 and other pneumonia patients, obtaining an overall accuracy of 90.2%, sensitivity of 89.9%, and specificity of 89.1%.

Chest radiograph acquisition is relatively simple with less radiation exposure than CT scans. However, a single CR image fails to incorporate details of infections in the lung and cannot provide a comprehensive view for the lung infection diagnosis. CT scan, on the other hand, is an alternative imaging modality that incorporates the detailed structure of the lung and infected areas. Unlike CR images, CT scans generate cross-sectional images (slices) to create a 3D representation of the body. Consequently, there has been a surge of interest on utilizing 2D and 3D CT images to identify COVID-19 infection. For instance, (Yang et al., 2020) proposed a DenseNet-based model to classify manually selected slices with COVID-19 manifestations and pulmonary parenchyma into COVID-19 and normal classes. The underlying study achieved an accuracy of 92% on the patient-level classification by averaging slice-level probabilities followed by a threshold of $0.8$ on the averaged values. Furthermore, the dataset used to train and test the model does not include other types of pneumonia. Identified Drawback 1: Such methods require manual selecting slices demonstrating infection to feed the model, which makes the overall process time-consuming and only partially-automated. To extract features from all CT slices, (Li et al., 2020) first segmented the lung regions using a U-net based segmentation method (Ronneberger et al., 2015), and then used them to fine-tune a ResNet50 model, which was pre-trained on natural images from the ImageNet dataset (Deng et al., 2009). Extracted features are then combined using a max-pooling operation followed by a fully connected layer to generate probability scores for each disease type. Their proposed model achieved sensitivities of 90%, 87%, and 94% for COVID-19, Community Acquired Pneumonia (CAP), and non-pneumonia cases respectively. Identified Drawback 2: Such methods combine extracted features from all slices of a patient, with or without infection, which potentially results in lower accuracy as there are numerous slices without evidence of infection in a volumetric CT scan of an infected patient.

In the study by (Hu et al., 2020), segmented lungs are fed into a multi-scale CNN-based classification model, which utilizes intermediate CNN layers to obtain classification scores, and aggregates scores generated by intermediate layers to make the final prediction. Their proposed method achieves an overall accuracy of 87.4% in the three-way classification. (Zhang et al., 2020) proposed a two-stage method consisting of a Deeplabv3-based lung-lesion segmentation model (Yasuda and Takahumi Majima, 2017) followed by a 3D ResNet18 classification model (Hara et al., 2017) to identify lung lesions and abnormalities and use them to classify patients into COVID-19, community acquired pneumonia, and normal findings. They manually annotated chest CT scans into seven regions to train their lung segmentation model, which is a time-consuming and sophisticated task requiring high level of thoracic radiology expertise to accomplish. Their proposed method achieves the overall accuracy of 92.49% in both three-way and binary (COVID-19 versus others) classifications.

At one hand, we aim to address the two identified drawbacks of the aforementioned methods. More specifically, existing solutions either require a precise annotation/labeling of lung images, which is time-consuming and error-prone, especially when we are facing a new and unknown type of disease such as COVID-19, or assign the patient-level label to all the slices. On the other hand, CNN, which is widely adopted in COVID-19 studies, suffers from an important drawback that reduces its reliability in clinical practice. CNNs are required to be trained on different variations of the same object to fully capture the spatial relations and patterns. In other words, CNNs, commonly, fail to recognize an object when it is rotated or transformed. In practice, extensive data augmentation and/or adoption of huge data resources are needed to compensate for the lack of spatial interpretation. As COVID-19 is a relatively new phenomenon, large datasets are not easily accessible, especially due to strict privacy preserving constraints. Furthermore, most COVID-19 cases have been reported with a specific infection distribution in their image (Bai et al., 2020; Chung et al., 2020; Shi et al., 2020; Ng et al., 2020), which makes capturing spatial relations in the image highly important.

As stated previously, structure of infection spread in the lung for COVID-19 is not yet fully understood given its recent and abrupt emergence. Furthermore, COVID-19 has a particular structure in affecting the lung, therefore, picking up those spatial structures are significantly important. Capsule Networks (CapsNets) (Hinton et al., 2018), in contrast to CNNs, are equipped with routing by agreement process enabling them to capture such spatial patterns. Even without a large dataset, capsules interpret the object instantiation parameters, besides its existence, and by reaching a mutual agreement, higher-level objects are developed from lower-level ones. The superiority of Capsule Networks over their counterparts has been shown in different medial image processing problems (Afshar et al., 2018, 2019a, 2019b, 2020c, 2020e, 2020d). Recently, we proposed a Capsule Network-based framework (Afshar et al., 2020b), referred to as the COVID-CAPS, to identify COVID-19 cases from chest radiographs, which achieved an accuracy of 98.3%, a specificity of 98.6%, and a sensitivity of 80%. As stated previously, CT imaging is superior for COVID-19 detection and diagnosis purposes when compared to chest radiographs. However, as in the case of CT imaging, we are dealing with 3D inputs and several slices per patient (compared to one chest radiograph per patient), the learning process is significantly more challenging. As such, accuracies of deep models trained over CT scans cannot be directly compared with those obtained based on chest radiographs.

Following our previous study on chest radiographs, in the present study, we take one step forward and propose a fully automated two-stage Capsule Network-based framework, referred to as the COVID-FACT, to identify COVID-19 patients using chest CT images. Based on our in-house dataset, COVID-FACT achieves an accuracy of 90.82%, sensitivity of 94.55%, specificity of 86.04%, and Area Under the Curve (AUC) of 0.98. We developed two variants of the COVID-FACT, one of which is fed with the whole chest CT image, while the other one utilizes the segmented lung area as the input. In the latter case, instead of using an original chest CT image, first a segmentation model (Hofmanninger et al., 2020) is applied to extract the lung region, which is then provided as input to the COVID-FACT. This will be further clarified in Section 5. Experimental results show that the model coupled with lung segmentation achieves the same overall accuracy compared to the other COVID-FACT variation working with original images. However, using the segmented lung regions increases the sensitivity and AUC from 92.72% and 0.95 to 94.55% and 0.98, respectively, while slightly decreasing the specificity from 88.37% to 86.04%. The $95\%$ Confidence Interval (CI) is also provided for all performance matrics using the methods described in (Brown et al., 2001; Hanley and McNeil, 1982).

COVID-FACT benefits from a two-stage design, which is of paramount importance in COVID-19 detection using CT scans, as a CT examination is typically associated with hundreds of slices that cannot be analyzed at once. At the first stage, the proposed COVID-FACT detects slices demonstrating infection in a 3D volumetric CT scan to be analyzed and classified at the next stage. At the second stage, candidate slices detected at the previous stage are classified into COVID and non-COVID (community acquired pneumonia and normal) cases and a voting mechanism is applied to generate the classification scores in the patient level. COVID-FACT’s two-stage architecture has the advantage of being trained by even weakly labeled dataset, as errors at the first stage can be compensated at the second stage. As a result, COVID-FACT does not require any infection annotation or a very precise slice labeling, which is a valuable asset due to the limited knowledge and experience on the novel COVID-19 disease. The trained COVID-FACT model and the code to implement the corresponding fully automated framework are available publicly for open access at https://github.com/ShahinSHH/COVID-FACT.

The reminder of the paper is organized as follows: Section 4 describes the dataset and imaging protocol used in this study. Section 5 presents a brief description of Capsule Networks and explains the proposed COVID-FACT in details. Experimental results and model evaluation are presented in Section 6. Finally, Section 7 concludes the work.

This research work is performed based on the policy certification number $30013394$ of Ethical acceptability for secondary use of medical data approved by Concordia University. The dataset used in this study, referred to as the “COVID-CT-MD” (Afshar et al., 2020a), contains volumetric chest CT scans of 171 patients positive for COVID-19 infection, 60 patients with community acquired pneumonia, and 76 normal patients acquired from April 2018 to May 2020. The average age of patients is $50\pm 16$ including 183 men and 124 women. Diagnosis of COVID-19 infection is based on positive real-time reverse transcription polymerase chain reaction (rRT-PCR) test results, clinical parameters, and CT scan manifestations by a thoracic radiologist, with 20 years of experience in thoracic imaging. Community acquired pneumonia and normal cases were included from another study and the diagnosis was confirmed using clinical parameters, and CT scans. A subset of 55 COVID-19, and 25 community acquired pneumonia cases were analyzed by the radiologist to identify and label slices with evidence of infection as shown in Figure 1. This labeling process focuses more on distinctive manifestations rather than slices with minimal findings. The labeled subset of the data contains 4,962 number of slices demonstrating infection and 18,447 number of slices without infection. The data is then used to train and validate the first stage of our proposed COVID-FACT model for extracting slices demonstrating infection from a volumetric CT data to be used in the second classification stage. We divided this subset into three separate components for training, validation, and testing. 60% of the labeled data is used for training, 10% for validation, and 30% for the test. The unlabeled subset is divided with the same proportion and used along with the labeled data to develop the second stage of the COVID-FACT model and evaluate the overall method.

All CT examinations have been acquired using a single CT scanner with the same acquisition setting and technical parameters, which are presented in Table 1, where kVP (kiloVoltage Peak) and Exposure Time affect the radiation exposure dose, while Slice Thickness and Reconstruction Matrix represent the axial resolution and output size of the images, respectively Raman et al. (2013). Next, we describe the proposed COVID-FACT framework followed by the experimental results.

A Capsule Network (CapsNet) is an alternative architecture for CNNs with the advantage of capturing hierarchical and spatial relations between image instances. Each Capsule layer utilizes several capsules to determine existence probability and pose of image instances using an instantiation vector. The length of the vector represents the existence probability and the orientation determines the pose. Each Capsule $i$ consists of a set of neurons, which collectively create the instantiation vector $\bm{u}_{i}$ for the associated instance. Capsules in lower layers try to predict the output of Capsules in higher levels using a trainable weight matrix $\bm{W}_{ij}$ as follows

$$\hat{\bm{u}}_{j|i}=\bm{W}_{ij}\bm{u}_{i},$$

(1)

where $\hat{\bm{u}}_{j|i}$ is the predicted output of Capsule $j$ in the next layer by the Capsule $i$ in the lower layer. The association between the prediction $\hat{\bm{u}}_{j|i}$ and the actual output of Capsule $j$, denoted by $\bm{v}_{j}$, is determined by taking the inner product of $\hat{\bm{u}}_{j|i}$ and $\bm{v}_{j}$. The higher the inner product, the more contribution of the lower level capsules to the higher level one. The contribution of Capsule $i$ to the output of the Capsule $j$ in the next layer is determined by a coupling coefficient $c_{ij}$, trained over a course of few iterations known as the “Routing by Agreement” given by

	$\displaystyle a_{ij}$	$\displaystyle=$	$\displaystyle\bm{v}_{j}.\hat{\bm{u}}_{j|i},$		(2)
	$\displaystyle b_{ij}$	$\displaystyle=$	$\displaystyle b_{ij}+a_{ij},$		(3)
	$\displaystyle c_{ij}$	$\displaystyle=$	$\displaystyle\frac{\exp(b_{ij})}{\sum_{k}\exp(b_{ik})},$		(4)
	$\displaystyle\bm{s}_{j}$	$\displaystyle=$	$\displaystyle\sum_{i}{c_{ij}\hat{\bm{u}}_{j|i}},$		(5)
	$\displaystyle\text{and}\qquad\bm{v}_{j}$	$\displaystyle=$	$\displaystyle\frac{\left\|\bm{s}_{j}\right\|^{2}}{1+\left\|\bm{s}_{j}\right\|^{2}}\frac{\bm{s}_{j}}{\left\|\bm{s}_{j}\right\|},$		(6)

where $a_{ij}$ is referred to as the agreement coefficient between the prediction and actual output, and $b_{ij}$ denotes the log prior of the coupling coefficient $c_{ij}$. Vector $\bm{s}_{j}$ denotes the Capsule output before applying the squashing function. As the length of output vectors represents probabilities, the ultimate output of Capsule $j$ ($\bm{v}_{j}$) is obtained by squashing $\bm{s}_{j}$ between 0 and 1 using the squashing function defined in Eq. (6). In order to update weight matrix $\bm{W}_{ij}$ through a backward training process, the loss function is calculated for each Capsule $k$ as follows

$$l_{k}=T_{k}\max(0,m^{+}-||\bm{v}_{k}||)^{2}+\lambda(1-T_{k})\max(0,||\bm{v}_{k}||-m^{-})^{2},$$

(7)

where $T_{k}$ is 1 when the class $k$ is present and 0 otherwise. $m^{+}$, $m^{-}$, and $\lambda$ are hyper parameters of the model and are originally set to 0.9, 0.1, and 0.5, respectively. The overall loss is the summation of all losses calculated for all Capsules.

The overall architecture of the COVID-FACT is illustrated in Figure 2, which consists of a lung segmentation model at the beginning followed by two Capsule Network-based models and an average voting mechanism coupled with a thresholding approach to generate patient-level classification results. The three components of the COVID-FACT are as follows:

• •

  Lung Segmentation: The input of the COVID-FACT is the segmented lung regions identified by a U-net based segmentation model (Hofmanninger et al., 2020), referred to as the “U-net (R231CovidWeb)”, which has been initially trained on a large and diverse dataset including multiple pulmonary diseases, and fine-tuned on a small dataset of COVID-19 images. The Input of the U-net (R231CovidWeb) model is a single slice with the size of $512\times 512$. The output is the lung tissues, which will be normalized between 0 and 1 to generalize the features and help the model to perform more effectively. Following the literature (Zhang et al., 2020; Hu et al., 2020), we down-sampled the output from $512\times 512$ to $256\times 256$ size to reduce the complexity and memory allocation without losing any significant information. Slices with no detected lung regions are removed and the remaining are fed into the first stage of the COVID-FACT model.

• •

  COVID-FACT’s Stage One: The first stage of the COVID-FACT, shown in Figure 3 is responsible to identify slices demonstrating infection (by COVID-19 or other types of pneumonia). Using this stage, we discard slices without infection and focus only on the ones with infection. Intuitively speaking, this process is similar in nature to the way that radiologists analyze a CT scan. When radiologists review a CT scan containing numerous consecutive cross-sectional slices of the body, they identify the slices with an abnormality in the first step, and analyze the abnormal ones to diagnose the disease in the next step. Existing CT-based deep learning processing methods either use all slices as a 3D input to a classifier, or classify individual slices and transform slice-level predictions to the patient-level ones using a threshold on the entire slices (Rahimzadeh et al., 2020). Determining a threshold on the number or percentage of slices demonstrating infection over the entire slices is not precise, as most pulmonary infections have different stages with involvement of different lung regions (Yu et al., 2020). Furthermore, a CT scan may contain different number of slices depending on the acquisition setting, which makes it impossible to find such a threshold. In most methods passing all slices as a 3D input to the model, the input size is fixed and the model is trained to assign higher scores to slices demonstrating infection. However, the performance of such models will be reduced when testing on a dataset other than the dataset on which they are originally trained (Zhang et al., 2020).

  The model used in stage one of the proposed COVID-FACT is adapted from the COVID-CAPS model presented in our previous work (Afshar et al., 2020b), which was developed to identify COVID cases from chest radiographs. The first stage consists of four convolutional layers and three capsule layers. The first and second layers are convolutional ones followed by a batch-normalization. Similarly, the third and fourth layers are convolutional ones followed by a max-pooling layer. The fourth layer, referred to as the primary Capsule layer, is reshaped to form the desired primary capsules. Afterwards, three capsule layers perform sequential routing processes. Finally, the last Capsule layer represents two classes of infected and non-infected slices. The input of stage one is set of CT slices corresponding to a patient, and the output is slices of the volumetric CT scan demonstrating infection. The output of stage one may vary in size for each patient due to different areas of lung involvement and phase of infection.

  In order to cope with our imbalanced training dataset, we modified the loss function, so that a higher penalty rate is given to the false positive (infected slices) cases. The loss function is modified as follows

  $$loss=\frac{N^{+}}{N^{+}+N^{-}}\times loss^{-}+\frac{N^{-}}{N^{+}+N^{-}}\times loss^{+},$$

  (8)

  where $N^{+}$ is the number of positive samples, $N^{-}$ is the number of negative samples, $loss^{+}$ denotes the loss associated with positive samples, and $loss^{-}$ denotes the loss associated with negative samples.

  

• •

  COVID-FACT’s Stage Two: As mentioned earlier, we need to apply classification methods on a subset of slices demonstrating infection rather than on the entire slices in a CT scan. It is worth noting that, lung segmentation (i.e., extracting lung tissues) is performed in one of the variants of the COVID-FACT as a pre-processing step. The first stage of the COVID-FACT, on the other hand, is tasked with this specific issue of extracting slices demonstrating infections.

  The second stage of the COVID-FACT takes candidate slices of a patient detected in stage one as the input, and classifies them into one of COVID-19 or non-COVID (including normal and pneumonia) classes, i.e., we consider a binary classification problem. Stage two is a stack of four convolutional and two capsule layers shown in Figure 4. The output of the last capsule indicates classification probabilities in the slice-level. An average voting function is applied to the classification probabilities, in order to aggregate slice-level values and find the patient-level predictions as follows

  $$P(p_{k}\in c)=\frac{1}{L_{k}}\sum_{i=1}^{L_{k}}P(s_{i}^{k}\in c),$$

  (9)

  where $P(p_{k}\in c)$ refers to the probability that patient $k$ belongs to the target class $c$ (e.g., COVID-19), $L_{k}$ is the total number of slices detected in stage one for patient $k$, and $P(s_{i}^{k}\in c)$ refers to the probability that the $i^{\text{th}}$ slice detected for patient $k$ belongs to the target class $c$. It is worth noting that while, initially, the COVID-FACT performs slice-level classification in its second stage, the output is patient-level classification (through its voting mechanism), which is on par with other works that COVID-FACT is compared with. As a final note to our discussion, we would like to add that, corona virus infection is, typically, distributed across the lung volume as such manifests itself in several CT slices. Therefore, having a single slice identified as COVID-19 infection can not necessarily lead to a positive COVID-19 detection.

  Similar to stage one, the loss function modification in Eq. (8) is used in the training phase of Stage two. The default cut-off probability of 0.5 is chosen in Stage two to distinguish COVID-19 and non-COVID cases. However, it is worth mentioning that the main concern in the clinical practice is to have a high sensitivity in identifying COVID-19 positive patients, even if the specificity is not very high. As such, the classification cut-off probability can be modified by physicians using the ROC curve shown in Figure 6 in order to provide a desired balance between the sensitivity and the specificity (e.g., having a high sensitivity while the specificity is also satisfying). In other words, physicians can decide how much certainty is required to consider a CT scan as a COVID-19 positive case. By choosing a cut-off value higher than 0.5, we can exclude those community acquired pneumonia cases that contain highly overlapped features with COVID-19 cases. On the other hand, by selecting a lower cut-off value, we will allow more cases to be identified as a COVID-19 case.

  To further improve the ability of the proposed COVID-FACT model to distinguish COVID-19 and non-COVID cases and attenuate effects of errors in the first stage, we classify all patients with less than 3% of slices demonstrating infection in the entire volume as a non-COVID case. These cases are more likely normal cases without any slices with infection. The few slices with infection identified for these cases might be due to the model error in the first stage, non-infectious abnormalities such as pulmonary fibrosis, or motion artifacts in the original images, which will be covered by this threshold. Based on (Yu et al., 2020), it can be interpreted that 4% lung involvement is the minimum percentage for COVID-19 positive cases. In addition, the minimum percentage of slices demonstrating infection detected by the radiologist in our dataset is 7%, and therefore 3% would be a safe threshold to prevent mis-classifying infected cases as normal.

• •

  Grad-CAM: Using the Grad-CAM approach, we can visually verify the relation between the model’s prediction and the features extracted by the intermediate convolutional layers, which ultimately leads to a higher level of interpretability of the model. Grad-CAM’s outcome is a weighted average of the feature maps of a convolutional layer, followed by a Rectified Linear Unit (ReLU) activation function, i.e.,

  $$L_{Grad-CAM}^{c}=RelU\left(\sum_{k}\alpha_{k}^{c}A^{k}\right),$$

  (10)

  where $L_{Grad-CAM}^{c}$ refers to the Grad-CAM’s output for the target class $c$; $\alpha_{k}^{c}$ is the importance weight for the feature map $k$ and the target class $c$, and; $A^{k}$ refers to the feature map $k$ of a convolutional layer. The weights $\alpha_{k}^{c}$ are obtained based on the gradients of the probability score of the target class with respect to an intermediate convolutional layer followed by a global average pooling function as follows

  $$\alpha_{k}^{c}=\frac{1}{Z}\sum_{i}\sum_{j}\frac{\partial y^{c}}{\partial A_{ij}^{k}},$$

  (11)

  where $y^{c}$ is the prediction value (probability) for target class $c$, and $Z$ refers to the total number of feature maps in the convolutional layer.