Different Deep Learning Based Classification Models for COVID-19 CT-Scans and Lesion Segmentation Through the cGAN-UNet Hybrid Method

In the past few years, machine learning has become one of the most popular and accurate methods for automated medical image analysis [17]. Deep Learning is a subfield of machine learning concerned with algorithms inspired by the structure and function of the brain called “artificial neural networks” (ANNs). Convolutional Neural Networks (CNNs) are deep learning algorithms that have mostly been used recently because they can automatically learn feature representation from a given image hierarchically. They can learn low-level to high level representations of an image and can then assign importance (learnable weights and biases) to various aspects/objects in the image and be able to differentiate one from the other. In this section, recent research that mainly focused on COVID-19 segmentation and classification in medical imaging are given.

2.1 Research on COVID-19 Segmentation

The goal of medical image segmentation is to identify ROIs such as organs or medical abnormalities by labeling each pixel in a medical image. Segmentation is a crucial step especially in COVID-19 image analysis as it can aid radiologists in diagnosis, disease progression monitoring, and improvements in speed and efficiency. However, for deep learning models to generalize well, they should be trained using large amounts of data and in certain tasks like image segmentation, the data must be well-labeled data.

The highly contagious nature of COVID-19 makes the data collection process particularly challenging as medical staff have to be well protected and equipment used in the process have to be frequently disinfected [21]. In addition, some datasets are not public. Despite these challenges, several approaches have been proposed for COVID-19 segmentation.

Due to the lack of annotated medical images, some researches involve unsupervised and semi-supervised methods. Zheng et al. [23] presented an unsupervised learning technique that generates pseudo segmentation masks from COVID-19 CT scans. By equipping the convolutional blocks with the so-called bottleneck blocks, Shan et al. [24] used a VB-Net to segment COVID-19 infection regions in CT scans. Fan et al. [25] proposed a semi-supervised approach called Inf-Net for identification of infected regions in CT slices.

The encoder-decoder architecture is commonly used for semantic segmentation [26]. In encoder-decoder architecture, the encoder learns the feature representations of the input while the decoder takes this feature representation, recovers the location information that was lost during the pooling process and produces a binary mask. One of such well-known architectures is UNet which retains the vital information from input images by adding skip connections between the encoding and decoding layers [27]. UNet and its variants have been developed for the segmentation of COVID-19. The 3D UNet [28] replaces the layers of the conventional UNet with a 3D version, thereby using the inter-slice information for segmentation. Zhou et al. [29] proposed a UNet based segmentation network that incorporates the attention mechanism to extract useful features from the encoders. They used spatial and channel attention to re-weight the feature representation and capture rich contextual relationships. But the result dice score was only 83.1%. After, Zhou et al. [30] proposed UNet++ where they modify UNet by inserting a nested convolutional structure between the encoding and decoding path. This network is adopted in Chen et al. [31] for lesion segmentation. They introduced a deep learning algorithm for automated segmentation of multiple COVID-19 infection regions. In order to learn a robust and expressive feature representation they used aggregated residual transformations and applied soft attention mechanisms. After segmenting the lesion, a 3-consecutive slice and quadrant based post-processing methods was applied to determine positive and negative scans. The obvious drawback to UNet architectures is that learning may slow down in the middle layers of deeper models, so there is some risk of the network learning to ignore the layer.

Wu et al. [32] developed a Joint Classification and Segmentation (JCS) system that performs real-time and explainable COVID-19 diagnosis. But their model achieved only a dice of 78% on a large scale dataset of 144167 CT images which the authors created. 

Another popular architecture for semantic segmentation is the Fully Convolutional Network (FCN) [33] where the fully connected layers of the CNN are removed. Milletari et al. [34] used a FCN with residual blocks as the basic convolutional block. Their proposed architecture called V-Net is made for 3D volumes and Dice loss is used to optimize the network. FCN classifies each pixel in a spliced image as spliced or authentic. But it often loses or smooths detailed structures and ignores small objects. In turn, Oktay et al. [35] brought forward an Attention UNet that uses advanced attention mechanisms to capture fine structures in COVID-19 images. The network was capable of segmenting particularly variable small size organs such as the pancreas.

Several approaches using Generative adversarial networks (GANs) also have been proposed for image segmentation. However, these approaches are based on augmenting the dataset by synthesizing COVID-19 CT scans. Zhang et al. [36] proposed CoSinGAN which can be learned from a single radiological image with the annotation mask of infected regions. The model was able to synthesize high-resolution and diverse images that match the input conditions.

2.2 Research on COVID-19 classification

Medical image classification aims to label a complete image to predefined classes and can therefore be used for diagnosis. Many approaches have been proposed in literature to classify COVID-19. Zhang et al. [37] propose a ResNet based model that performs anomaly detection and classification between COVID-19 and non-COVID-19 X-ray images. The anomaly score is used to optimize classification on the dataset that contains 10078 images. The model produces a sensitivity, specificity and AUC score of 96.0%, 70.7% and 95.2%, respectively.

Loey et al. [38] implemented a GAN with deep transfer learning for the detection of coronavirus. They used the GAN to generate more images from publicly available X-ray images. These images were then classified in the first scenario using AlexNet, ResNet18 and GoogleNet which produce accuracies of 66.67%, 69.67% and 80.56% respectively. Apostolopoulos and Mpesiana [39] used transfer learning for the detection of COVID-19 with Vgg19 obtaining the best accuracy of 96.78%. After, Yang et al. [40] built an open-source dataset and showed that a CNN model can be used for diagnosing COVID-19. Their proposed model achieved an accuracy and F1-score of 89% and 90%, respectively. Hemdan et al. [41] proposed CovidX-Net which includes seven different transfer learning architectures for the diagnosis of COVID-19. Vgg19 and DenseNet achieved the best performance with 90% accuracy.

Sethy et al. [42] used support vector machines (SVM) to classify features obtained from several convolutional neural networks. In their study, the ResNet50 model with SVM classifier achieved best results with an accuracy of 95.33%. Similarly, Song et al. [43] utilized a Feature Pyramid Network (FPN) and attention modules to extract features from CT images and fed the extracted features into a ResNet-50 model. They obtained an accuracy of 86%. Ozturk et al. [44] presented DarkCovidNet which automatically detects COVID-19 from X-ray images with an accuracy of 98.08%. DarkCovidNet was inspired from Darknet-19 which is the backbone of a state-of-the-art real time object detection system YOLO (You Only Look Once). Covid-Net which proposed by Wang et al. [8] was trained on a large dataset of 13,975 X-ray images. The model architecture makes used of projection-expansion-projection design pattern which provides enhanced representational capacity. 

Wang et al. [45] proposed a 2D CNN model that performs classification between COVID-19 and viral Pneumonia. Experimental results showed that the proposed model achieves an accuracy of 73.1%, sensitivity of 74.0% and specificity of 67.0%. Covid-Net, a deep CNN based model was proposed by Wang et al. [8] to detect COVID-19 in X-ray images producing a testing accuracy of 83.5%. Narin et al. [46] on the other hand, employed several transfer learning methods to detect COVID-19. Specifically they used ResNet50, InceptionV3 and an ensemble Inception-ResNetV2 on two different datasets. While ResNet50 attained the best performance with an accuracy of 98%, InceptionV3 and Inception-ResNetV2 achieved accuracies of 97% and 87%, respectively. Additionally, Ghoshal and Tucker [47] presented a Bayesian Convolutional Neural network that improves the classification accuracy of the standard VGG16 model from 85.7% to 92.8%. The authors enhanced explainability by using saliency maps to point out the specific locations that the model focuses on to make its decisions.

Li et al. [48] introduced Covid-MobileXpert which is a lightweight deep neural network based mobile app that can be used for the screening of COVID-19. The model consists of a 3-player knowledge transfer and distillation (KTD) framework that includes a pre-trained attending physician network (DenseNet-121), a resident fellow network and a medical student network. Then, Lahsaini et al. [49] presented a transferred DenseNet-201 model for classification of COVID-19 from CT images that give the classification accuracy of 98.18%. Bargshady et al. [50] used CycleGAN to augment the dataset and then classified the images with a modified Inception model with 94.2% accuracy. 

Several proposed methods performed both segmentation and classification. Zheng et al. [32] employed a U-Net model for lung segmentation and the result is used as an input to a 3D CNN to predict the presence of COVID-19. A sensitivity of 90.7%, specificity of 91.1% and AUC of 95.9% was achieved with this model. Similarly, Chen et al. [28] used UNet++ to obtain segmented lesions from CT images and, to predict COVID-19 and other forms of viral pneumonia.  With the proposed model, the reading time of radiologists shortened by 65%. Xu et al. [51] used a V-Net to segment candidate infected regions and combined the resulting output with handcrafted features. This combination was then sent to a ResNet-18 network that attains a classification accuracy of 86.7%. Li et al. [52] used a combination of U-Net and ResNet50 to extract lung regions and performed classification with 96% specificity, 90% sensitivity and an AUC of 96%. Amyar et al. [53] proposed a multi-task deep learning model which jointly classifies, reconstructs and segments COVID-19 lesions from CT images. Their architecture comprised a common encoder, 2 decoders and a multi-layer perceptron. Their model produced a dice coefficient of 96% accuracy. Karakanis and Leontidis [54] proposed a light weight deep neural network that uses synthetic images generated from a GAN to detect COVID-19. An accuracy of 98.7% was obtained from the proposed model.

In medical image processing, both segmentation and classification studies can be performed. The purpose of segmentation is to identify Regions of Interest (ROI) such as organs or medical abnormalities by labeling each pixel in a medical image. On the other hand, medical image classification aims to label a complete image into predefined classes, and for this, it first divides the image into specified sub-domains. Thus, only the ROI is examined in medical images divided into sub-areas, and classification is made according to these areas. However, ROI areas cannot always be obtained in medical images. In such a case, first ROI should be segmented and then the classification algorithm should be applied. In the COVID-19 classification dataset (Dataset-2) used in the study, there are no ROI areas as masked image areas. For this reason, first of all, a successful segmentation algorithm based on deep learning was proposed by using Dataset-1 that has ground truth masked images. After the success of the proposed segmentation algorithm has been proven on Dataset-1, the resulting segmentation model has been used on the Dataset-2 to generate mask which are then used for classification. So, ROI of lung was segmented and According to the classification results, images could be classified more accurately in a short time compared to the literature. Since both segmentation and classification were made within the scope of the study, the literature studies and results in these two fields were examined separately and comparisons were made with the results.

Our proposed method can be summarized in two main stages (Figure 2) as segmentation and classification. The data was preprocessed before feeding it into the model.

After preprocessing, the data has been fed into a deep learning model to obtain a binary mask of lesions in COVID-19 CT images in Stage-1. For this purpose; UNet, UNet++, SegNet and PsPNet were used both separately and as hybrids with GAN, to automatically segment infected areas from chest CT slices. The resulting binary mask was then fed into Stage-2, where a CNN, PatchCNN and CapsNet models were used to predict whether or not a patient has COVID-19 or not. The binary masks that contain COVID-19 have white areas on a black background while those without COVID-19 are completely black. Steps of the algorithm are given in Figure 3.

2.png

Figure 2. Steps of the proposed method

3.png

Figure 3. Steps of the algorithm

4.1 Preprocessing

Before the first stage, several preprocessing steps are applied to the input images. First, 2D slices are extracted from the provided CT scans. Then, images are converted to gray-scaled and normalized. Generally, medical image analysis is challenging due to inherent characteristics of medical images like low contrast, noise, signal dropouts and complex anatomical structures [57]. In order to solve this problem, Contrast Limited Adaptive Histogram Equalization (CLAHE) is applied to the CT scans. It enhances small details, textures and the local contrast of the images [58]. CLAHE is a variant of adaptive histogram equalization where the contrast amplification is limited to reduce noise amplification. This is done by clipping the value of a histogram at each pixel value and then redistributing it equally among all histogram bins [59].

Most of the images in the dataset contained a lot of black space and sections below the lungs which could lead to wastage of computing resources. The proposed solution to this problem was extracting the ROI by drawing contours over the image and then cropping out the rectangle with the largest closed boundary. The resulting ROI encompasses both lungs in the image. Furthermore, the ROI images are resized to 256 x 256 pixels and then normalized to pixel values between [0,1] using min-max scaling. To avoid overfitting and increase the number of training images data augmentation strategies like rotation, shearing, zooming and horizontal flipping were applied to the training data. The corresponding segmentation masks were also cropped, resized, normalized and augmented. Figure 4 shows the outputs of the preprocessing stage.

4.2 Segmentation step

Image segmentation is an image processing problem that aims to grouping the pixels according to the similarities or dissimilarities (differences). Segmentation makes the image more meaningful and provides faster further stages [60].

Conditional Generative Adversarial Networks (cGAN) was used to segment COVID-19 lesions from CT images in this study. The generative network learns to recognize our ROI and creates a binary mask that outlines its boundaries. The discriminator on the other hand, learns to differentiate between the real segmented masks and the synthetic ones.

In this section, the overall architecture of the GAN, cGAN, the generator, discriminator and how they work are presented. The proposed architecture was inspired by the Pix-2-Pix GAN [61] which is used for image to image translation.

4.png

Figure 4. Steps in the order of pre-processing

4.2.1 Generative Adversarial Network (GAN) and Conditional GAN (cGAN)

A GAN is composed of a generator and a discriminator that are trained to be competitive with each other. Vanilla GAN that is basic type of the GAN can be seen in Figure 5. The generator, G; is trained to produce realistic fake data G(z) from the distribution of the real data (x) and the random noise vector z, in order to fool or confuse the discriminator. The goal of the discriminator, D, on the other hand, is to distinguish between the real data and the fake data produced by the generator. The generator is updated only through gradients back-propagated from the discriminator. By iteratively updating both the generator and discriminator, equilibrium is reached where the discriminator is no longer able to tell the difference between the real and fake data produced by the generator. In Mathematical terms, D and G play a two-player minimax game with objective function given by Eq. (1) [62, 63].

5.png

Figure 5. Vanilla GAN [64]

In Eq. (1), L is the objective function, G is generator, D is discriminator, E is the expected value, x is the input image, P_data(x) is the distribution from which training examples are drawn and z is random noise. The discriminator D is trained to maximize the probability of correct label assignment to fake and real data while the generator G is trained to fool the discriminator by minimizing log(1-D(G(z))) (that is the log of 1 minus the discriminator output of the generated/fake data). Hence, the gradients that are back propagated through the discriminator are used to update the generator. Much stronger gradients can be obtained in earlier iterations, if the generator is optimized to maximize log(D(G(z))) instead of minimizing log(1-D(G(z))) [65]. This is an indirect way of optimizing and the advantage with this method is that it ensures that the generator doesn’t memorize the images. Even though having many advantages, GANs are known that training is challenging because they are generally unstable and sometimes fail to converge [66]. In addition, GANs depict a phenomenon called model collapse where the generator generates a limited diversity of samples, or even the same sample, regardless of the input. There are numerous variations of GANs that have been tried to overcome that limitations that vanilla GANs, and cGAN is one of these variations.

One limitation of GAN is that it has no control on the actual image generation. The cGAN proposed by Mirza and Osidero tries to overcome this problem by incorporating additional information like class labels in the image generation process [67]. In the cGAN the generator has 2 inputs; a random noise z and some prior information c. Then the prior information c is fed into the discriminator together with the corresponding real or fake data. The authors depicted that this change in architectural structure, cGAN not only increased the generation of detailed features, but also enhances stability during training. Because of this, cGAN is used in this study. Eq. (2) shows the objective function of the cGAN framework. Also, the difference in the architecture between GAN and cGAN is given in Figure 6.

$L_{G A N}(G, D)=\mathbb{E}_{x \sim P_{\text {data }(x)}} \quad [\log (D(x))]+\mathbb{E}_{x \sim p_{\text {data }(x)}} \quad [1-\log (D(G(z)))]$      (1)

$L_{c G A N}(G, D)=\mathbb{E}_{x \sim P_{\text {data }(x)}} \quad [\log (D(x \mid c))]+\mathbb{E}_{x \sim p_{z(z)}} \quad [1-\log (D(G(z \mid c)))]$     (2)

6.png

Figure 6. Difference of the architecture between GAN and cGAN [68]

4.2.2 Generator architecture

The generator of cGAN used in this study is a modified U-Net architecture hence the name (cGAN-UNet). U-Net which was introduced in [24] is an architecture that is popularly used in biomedical image segmentation. It is a “U” shaped architecture that is made up of 3 main sections as the encoder, the bottleneck and the decoder sections. The encoding section in the generator of the cGAN-UNet architecture follows typical convolutional neural network architecture and is made up of many encoder blocks. Each encoder block passes an input image through one convolutional layer, followed by batch normalization and a leaky rectified linear unit (slope 0.2). Unlike the original U-Net, all max pooling layers are removed. After each block, the number of kernels is doubled so that the architecture can learn complex structures effectively.

The bottleneck section connects the contraction and expansion layer. It is made up of one 4x4 convolutional layer with a 2x2 stride, followed by a ReLU activation layer. This is often referred to as skip connections. Just like the contraction section, the expansion section is made up of several blocks. In each decoder block the input image is passed through one deconvolutional layer, followed by batch normalization, a dropout layer, a concatenate layer and finally an activation layer. Unlike the encoder blocks, here, the number of feature maps is reduced by half in order to maintain symmetry. In the concatenate layer, the input to the corresponding contraction layer is appended. This makes sure that the features that were learnt during contraction are used to reconstruct the new image.  It should be noted that the number of encoder blocks is the same as the number of decoder blocks. The result from the expansion block is then passed through another deconvolutional layer where the number of feature maps is equal to the desired number of segments or classes. All convolutional and deconvolutional layers are defined with a 4x4 kernel and a stride of 2x2 which down sample and up sample the feature maps respectively. This architecture is displayed in Figure 7.

7.png

Figure 7. Generator architecture in the study

8.png

Figure 8. Size of feature vectors in input / output layers of the generator model

9.png

Figure 9. Discriminator architecture in the study

10.png

Figure 10. Size of feature vectors in input / output layers of the discriminator model

To improve the performance of cGAN-UNet, several strategies from [66-69] were applied in this study. These strategies were included as batch normalization, using leakyReLU, random Gaussian weight initialization with a standard deviation of 0.002 in both the discriminator and generator. In addition, L2 regularization was added to all layers of the discriminator. Furthermore, label smoothing was applied on all class labels. In label smoothing, rather than use hard labels like 1 and 0 to represent real and fake labels, soft labels were used by generating random values between 0.9 and 1.0 to represent real images and values between 0 and 0.1 to represent fake images. After implementing these changes, there was a noticeable change in how fast the model converged. Sizes of the all feature vectors both input and output layers of the generator is given in Figure 8.

4.2.3 Discriminator architecture

The discriminator determines if a given binary mask is real or fake. The discriminator used in this study is called a “patch discriminator” [61]. The patch discriminator divides the input image into a set of patches and maps each patch to a single scalar output. Unlike a normal image discriminator which performs prediction on an entire image, a patch discriminator makes a prediction for each patch and the final prediction is the average of all patch predictions. In addition, the patch discriminator requires fewer parameters, it works well with very large and blurry images and the computational time is lower. In the study, the patch discriminator uses a patch size of 70x70 pixels. The authors of [61] showed that a 70x70 patch size is a good choice as it can accommodate global features. The implemented model takes in 2 input images that are concatenated together. To prevent overfitting, Gaussian noise (standard deviation of 0.2) is added to this concatenated input before passing into the first hidden layer. The model consists of 3 hidden layers, each followed by batch normalization and leaky ReLU. Each convolutional layer has 4x4 filters and a stride of 2. In addition, the discriminator is regularized by constraining the norms of its gradients using L2 regularization. Therefore, an L2 regularizer is added to all convolutional layers. The discriminator architecture is given in Figure 9. Sizes of the all feature vectors both input and output layers of the discriminator is given in Figure 10.

The discriminator is optimized using a combination of the weighted binary cross entropy (bce) loss (Eq. (3)) and the dice loss (Eq. (4)). The discriminator loss is calculated using (Eq. (5)) and this loss back propagated through the generator. The cross-entropy loss is the most widely used loss for CNN classification [70].

The generator is updated using the discriminator loss and the L₁ loss (Eq. (6)) with a weighting of 1 to 100 in favor of the L₁ loss as recommended by the authors of [61]. The final objective function for the generator (with λ=100) is shown in Eq. (7).

$L_{B C E}=-\frac{1}{N} \sum_{i=1}^N y_i \cdot \log \left(p\left(y_i\right)\right)+\left(1-y_i\right) \cdot \log \left(1-p\left(y_i\right)\right)$      (3)

$L_{D I C E}=1-\frac{2 x(y \cap p(y))}{y+p(y)}$     (4)

$L_{D I S}=L_{B C E}+L_{D I C E}$     (5)

$L_{L 1}(G)=\mathbb{E}_{x, y, z}\|y-G(x, z)\|$     (6)

$L_{G E N}=\arg _G^{\min \max } L_{C G A N}(G, D)+\lambda L_{L 1}(G)$     (7)

where, y is the ground truth, p(y) is the predicted mask for all values of N that is the number of images in the dataset and L_DIS is the discriminator loss, λ is the constant used to increase the weight of L₁ loss, L_L1 is the L₁ loss, L_GEN is the generator loss and L_CGAN is the adversarial loss as shown in Eq. (2).

The generator and discriminator are optimized simultaneously where the generator learns to create a reasonable binary mask and the discriminator learns how to differentiate between synthetic and real segmentation.  Besides introducing the weighted bce dice loss, Gaussian noise and L2 regularization, different learning rates and loss functions were used. Adam is found to be the best optimizer with a learning rate of 0.0001 and momentum of β₁ of 0.5 and β₂ of 0.999. However, the model was trained for 130800 epochs to achieve best results.

It should be noted that after generating the segmentation masks, a simple thresholding was performed to generate binary masks.  54.4 was selected as the best threshold value at which the 92.23 F1-score value was obtained as shown in Figure 11.

11.png

Figure 11. Threshold to F1-score graph

Training parameters of the cGAN-UNet used in the study also given in Table 2.

4.3 Classification step

After generating segmentation maps, they are used as input to a several classifiers to predict whether the masks have COVID-19 or not. Before feeding the synthetic masks into the classifiers, simple preprocessing and data augmentation methods like rotation, shearing, zooming and Gaussian blurring was applied as shown in Figure 12. Three different classification methods are implemented and their performance compared with experimental results. The architectures of these methods are presented the following section.

Table 2. Training parameters for cGAN-UNet

4.3.1 Convolutional Neural Network (CNN)

CNN’s are popularly used for image classification as they accumulate sets of features (such as edges, corners, shapes etc) at each layer. As shown in Figure 13, the CNN architecture is made up of 5 convolutional layers and 3 dense layers. While each convolutional layer is followed by batch normalization and a max pooling layer, the dense layers are followed by dropout layers. The ReLU activation function is applied on all layers except the final layer that uses sigmoid as an activation function. The convolutional layers use 3x3 kernels and have a stride of 1. The network is trained for 150 epochs. Batch normalization and dropout layers are added to avoid overfitting. The feature vectors of the CNN are given in Figure 14 and training parameters in Table 3.

Table 3. Training parameters for CNN and PatchCNN

12.png

Figure 12. Preprocessing of synthetic mask

13.png

Figure 13. Convolutional neural networks

14.png

Figure 14. Size of feature vectors in input / output layers of the CNN and PatchCNN model

4.3.2 Patch Convolutional Neural Networks (PatchCNN)

The PatchCNN architecture is very similar to the CNN architecture. However, unlike like CNN where the entire image is sent into the network, several patches are extracted from an image then each patch is sent into the network. In the proposed method, the class label for each patch is predicted by the model, then a majority voting scheme is applied.

The class that receives the most votes is the final class of the image. Several patch sizes were tried on the synthetic images generated from the cGAN-UNet model whose size is 256 x 256 pixels. The patches are created by dividing the synthetic image into smaller images of a specific patch size. As can be seen in Figure 15, the patch sizes we experimented with are 16 x 16, 32 x 32, 64 x 64 and 128 x 128 pixels.

As shown in Table 4, the bigger the patch size, the better the model performance. This indicates that smaller patch sizes do not carry sufficient diagnostic information while larger patch sizes carry sufficient discriminating features. According to the patch size success, all PatchCNN models were trained with a patch size of 128. Hence, model input was 128 x 128 pixels.

The training parameters for this model are the same as the CNN as shown in Table 3 while the feature vectors can be seen in Figure 15. Figure 16 shows the model architecture and a synthetic image after it has been split into different patch sizes.

Table 4. Comparative analysis of different patch sizes on cGAN-UNet Synthetic Masks

4.3.3 Capsule Neural Network (CapsNet)

CNNs have several pooling layers which tend to lose information such as the objects precise position and pose. CapsNet was introduced by Sabour [71] in 2017 to overcome these drawbacks in CNNs using the idea of inverse graphics. Unlike CNNs, CapsNet preserves detailed information about an object’s position and pose throughout the network. A capsule network is made up of capsules; a capsule is a group of neurons that learn to detect an object or parts of an object in an image. Its output is a vector whose length represents the presence of an object in an image and whose orientation could represent the position, rotation, size or scale of the object in the image. This is unlike CNNs that ignore the spatial relationships between features.

15.png

Figure 15. Input preprocessed image and the different patches extracted from it

16.png

Figure 16. PatchCNN Architecture with a 128 x 128 pixel Patch size

The CapsNet architecture is made up of an encoder and a decoder. The encoder contains 3 layers which include Convolutional Layer, Primary Capsule Layer and ClassCaps Layer. The Convolutional Layer has 256 kernels with size 9x9x1 and stride of 1, followed by the ReLU activation function. In this layer, the image pixel intensities are converted to local feature detectors. The output of this layer is then fed to the Primary Capsule Layer which is also a convolutional layer that contains 32 primary capsules with each capsule having a dimension of 6x6. The lowest level entities in the image are captured in this layer. The output of this layer is then sent to the ClassCaps layer through a dynamic routing algorithm.  Dynamic routing algorithm computes a coupling coefficient to quantify the connection between the Primary Capsule and the ClassCaps layers while facilitating the routing of capsules to the appropriate next layers. The CapsNet  model learns through the coefficient value as it is updated during the training process and only 3 routing iterations are  used to optimize the loss faster and avoid overfitting. The ClassCaps layer has 2 class capsules, one for each class. An 8x16 weight matrix is applied to each input vector here that maps the input space to the 16 dimensional capsule output space. 

The purpose of the decoder network in the CapsNet is to reconstruct the vectors from the ClassCaps layer representation. It does this by minimizing the Euclidean distance between the input image and decoder’s output. The decoder is made up 3 fully-connected layers; the first layer has a total of 512 neurons, the second layer 1024 neurons and the third contains 784 neurons. The CapsNet parameters are optimized using a margin loss for each output class. In addition to the margin loss, an l2 reconstruction loss (mean squared error) is used as a regularization method to boost each capsule to encode parameters of the input image. The CapsNet architecture can be seen in Figure 17. A summary of CapsNet architecture and the parameters, used in this study are shown in Figure 18 and Table 5, respectively.

Table 5. Training Parameters for CapsNet

17.png

Figure 17. CapsNet Architecture [69]

18.png

Figure 18. Size of feature vectors in input / output layers of the CapsNet model

With the emergence of COVID-19 at the end of 2019, the world faced a massive health crisis. In this study, different hybrid deep learning methods for the segmentation and classification of COVID-19 were proposed [78]. First of all, a conditional GAN (cGAN) was used to create binary masks that could assist radiologists in diagnosing, monitoring disease progression, reduction in examination time, and improvement in accuracy. After that, to further ease the diagnosis process, several classification models were used as Convolutional Neural Network (CNN), a PatchCNN and a Capsule Neural Network (CapsNet) for COVID-19 detection. These classifiers take in the generated binary mask as input and predict the presence or absence of COVID-19.

Our proposed methods show very promising results as they outperform baseline and state of the art models. The proposed segmentation method achieved a dice score of 92.31% and an IOU score of 86.41%. For the classification results, the best model attained accuracy and an F1-score of 99.20% on an independent dataset that has no ground truth binary masks. Therefore, the proposed method has the potential to aid front line health care providers and can be deployed in areas with limited access to health care facilities.

Even though proposed segmentation model achieved a better performance than all other baseline models, the amount of time it took to train the model was long (about 72hrs). However, inference time for all models is approximately the same, approximately 2 seconds for one image. With higher computing power proposed method may produce faster results.

To prevent overfitting, several measures were used like batch normalization, dropout, stratified cross validation. In the future, a larger dataset can be acquired from hospitals to further validate results. Hence, the model’s extent to which the model generalizes in different locations can be evaluated. Additionally, a CapsNet model with many more layers may produce better results that the one used in the study.

Furthermore, applying better post-processing methods that can remove noisy small patches from generated masks may improve our results. Finally, the basic unit of the all models used in this study is the neural network which is a black-box model and its explainability is still a work in progress. Hence, if the model structures can be explained more clearly, different studies can be done in the future.