CapsNet-COVID19: Lung CT image classification method based on CapsNet model

1.   Introduction

1.1.   Research foundation

After the novel coronavirus pneumonia (Corona Virus Disease 2019, COVID-19) was first discovered in November 2019, it caused extremely high morbidity and mortality worldwide in a short period of time ^[1]. Radiological examination can be used for COVID-19 screening, that is, radiologists perform chest imaging (such as chest X-ray or computed tomography imaging) of the patient, and can look for visual indicators related to the COVID-19 virus infection ^[2]. Early studies have found that once COVID-19 patients are infected, there will be abnormalities in lung imaging ^[3,4,5]. For example, Huang et al. found in their research that most COVID-19 positive cases showed bilateral radiological abnormalities in CXR images. These abnormalities are most common with ground glass opacity, patchy shadows, warped spreading and lung consolidation. Therefore, researchers believe that lung imaging can be used as the main tool for COVID-19 screening in the hardest-hit areas ^[6].

In recent years, artificial intelligence (AI) technology represented by convolutional neural networks (CNN) has made great achievements in computer vision. As the CT radiography mode has been proved to have higher sensitivity and specificity than the RT-PCR detection method, many researchers have begun to use CT and CRX image data sets to train automatic classification AI algorithms. Ozturk et al. ^[7] proposed a deep model that uses DarkNet as a classifier and has an end-to-end architecture. However, the number of data sets needs to be increased to make the model more robust and accurate. Y. Pathak et al. ^[8] used transfer learning methods to classify patients infected with COVID-19. The author introduces the smooth loss function top-2 with cost-sensitive attributes to process noise and unbalanced data sets, and uses a 10-fold cross-validation method to prevent overfitting. Narin et al. ^[9] used three advanced deep learning models (ResNet50, Inception v3 and Inception-ResNetV2) to classify CRX images (50 COVID-19 images vs. 50 non-COVID-19 images). In June 2021, Adel Oulefki et al. proposed a method to automatically segment and measure COVID-19 lung infection area ^[10]. In November 2021, Chen Zhao et al. used chest CT imaging features for lung segmentation and COVID-19 automatic detection ^[11]. In December 2021, Juan Juan He et al. used the evolutionary countermeasure network with gradient punishment to realize the segmentation task of COVID-19 infection ^[12]. In December 2021, Nan Mu et al. applied progressive global perception and local polishing network to the segmentation task of pulmonary infection in COVID-19 CT image ^[13]. In 2021, Qimao proposed the intelligent immune clonal optimization algorithm for pulmonary nodule classification ^[14]. In February 2022, X. Liu et al. used scribble annotation to segment the infected area of COVID-19 CT image ^[15].

Main contributions:

(1) This research proposes a two-level deep network model, which mainly realizes the classification of lung CT images. The first level of the model uses convolutional neural networks to locate lung regions in lung CT images. The second level of the model uses the capsule network to classify and predict the segmented images.

(2) SegNet-Lung segmentation model is used to segment and preprocess the lung region of lung CT image (the segmentation model in this study is named SegNet-Lung model) in order to extract shape features such as color and texture from lung CT image. At the same time, the unlabeled lung CT images are predicted and labeled as the training/test data set in the classification processing stage.

(3) The CapsNet-COVID19 capsule network classification model is established to classify and reconstruct the segmented lung CT images, and the column diagram is used to visualize the prediction results. By stopping early and increasing the amount of training data, the over fitting of the model caused by excessive iterative training in the training process is avoided. The confusion matrix, accuracy, boundary loss, reconstruction loss and total loss are used to quantitatively evaluate the capsule network model.

(4) Our classification method is suitable for medical image classification with complex background, low recognition rate, blurred boundaries and large image noise. The accuracy of our method is 84.291% on the test set and 100% on the training set. We believe that this classification method is of great value for monitoring and controlling the growth of patients in COVID-19 infected areas.

1.2.   Introduction to data source and research process

All human lung CT image data used in this research are derived from the 2019 Novel Coronavirus Information Database (2019nCoVR) provided on the website of China National Bioinformatics Center ^[16]. The lung CT image data in this database is provided by the China Chest CT Image Investigation Association (CC-CCII), so 2019nCoVR has high professional and practical value. There are three main types of lung CT images in 2019nCoVR: The Novel coronavirus pneumonia data set caused by the SARS-CoV-2 virus, the common pneumonia data set, and the normal lung data set. The amount of data in 2019nCoVR is relatively large. In view of the limitations of the experimental hardware equipment in this research, we selected 29,278 images with clear and obvious lung texture characteristics from the information database as the data set of this research. An example of lung CT images in the dataset is shown in Figure 1.

The main features of lung CT images of COVID-19 patients are early unilateral or bilateral lungs with localized inflammatory infiltration, which is mainly manifested as patches, masses, segments or subpleural subsegmental ground glass shadow. In the late stage, consolidation is dominant in the lungs ^[17]. Common pneumonia is generally divided into lobar pneumonia and lobular pneumonia. The early stage of lobar pneumonia is mainly manifested as a ground-glass-like density shadow in a certain lung lobe with fuzzy edges. Lobular pneumonia is common in the middle and lower parts of the lungs. The local bronchial vascular bundles are thickened, and there are blurry patches of varying sizes. The texture of the lungs of normal people is very clear, and there is no thickening, disorder or abnormal lines. At the same time, there are no obvious high-density or low-density dense shadows and nodules in both lungs, and no pathological manifestations such as hyperplasia, cavitation, calcification and exudation. The CT images of common pneumonia and COVID-19 overlap to a certain extent, and it is often difficult to make an accurate judgment of these two types of pneumonia by the human eye alone. Therefore, it is necessary to increase the diagnostic accuracy of COVID-19 with the help of image analysis technology.

This research uses a deep network model to analyze lung CT images to achieve automatic segmentation and classification of a large number of lung CT images (three main types). First, use an image segmentation model similar to U-NET ^[18] to segment the lung area of the lung CT image (the segmentation model in this research is named SegNet-Lung segmentation model). This model can extract shape features such as color and texture from lung CT images. Then, a capsule network model ^[19] was established to classify and predict the segmented lung CT images (the classification model in this research is named CapsNet-COVID19 classification model). The specific process is as follows:

(1) Train the segmentation model SegNet-Lung segmentation model using the labeled lung CT image set. Then, the trained SegNet-Lung segmentation model is used to segment the lung CT images to be classified.

(2) Use a self-editing program to unify the format of the segmented lung CT images: (a) The edge of the lung region in the image predicted by SegNet-Lung segmentation model is outlined with red lines; (b) Extract the lung region of the outlined lung CT image and hide other irrelevant regions; (c) Convert the segmented lung region image format into the input format available for the CapsNet-COVID19 classification model.

(3) The above-mentioned image set is divided into train set, validation set and test set according to an appropriate ratio. Through training the CapsNet-COVID19 classification model, the prediction of image classification is finally realized.

2.   SegNet-Lung segmentation model and lung CT image segmentation

2.1.   Data preparation

The data sets of lung CT images used in this research include: COVID-19, common pneumonia (CP) and normal lungs (Normal). The numbers of them are 9420, 9858 and 10,000 respectively. Set the resolution of the image set to 128 × 128 and the format to.tif. If the lung CT images are directly classified, the image classification accuracy will be reduced because the images are interfered by the external environment when they are taken. Therefore, it is necessary to segment the lung area of the lung CT image. We have about 2532 images of COVID-19 lung CT masks that have been annotated by experts (as shown in Figures 3 and 4). Then, they are used as the train set and validation set of the SegNet-Lung segmentation model, and the unlabeled lung CT images to be classified are used as the test set. Among them, the training set and validation set (90%) are used to train the SegNet-Lung segmentation model to determine the network parameters of the segmentation model, and the test set (10%) is used to test and predict the recognition rate and generalization of the trained segmentation model (Table 1). The test set is used to test and predict the segmentation performance of the trained SegNet-Lung segmentation model. Finally, the trained SegNet-Lung segmentation model is used for segmentation prediction and labeling of unlabeled lung CT images. Therefore, the labeled lung CT images can be used as the training/test data set of CapsNet-COVID19 classification model in the next step.

2.2.   SegNet-Lung segmentation model research and experiment

The SegNet-Lung segmentation model used in this study is improved on the basis of the U-NET model. This model uses Nconvolution instead of MaxPool and adds a random dropout layer with a parameter of 0.5. At the same time, the Sigmoid function is used in the output convolutional layer. During model training, the batch_size parameter is set to 64, and the convolutional layer uses a 3 × 3 filter and activation function ELU, which can reduce the feature dimension of the subsequent pooling layer and avoid overfitting. The filter sizes used by the largest pooling layer and upsampling layer are 3 × 3 and 2 × 2. After the convolutional layers are merged, the network structure of the model will be flattened. The SegNet-Lung segmentation model structure is shown in Figure 5.

The structure of nconvolution of covseg net model is convolution + batchnormalization + ELU. Batch normalization (BN) is used in the convolution process, which can forcibly pull the more and more biased data distribution in model training back to the standard distribution through certain standardized means. In this way, the active input value can fall in the area where the nonlinear function is sensitive to the input, and the output of the network will not be very large. Therefore, using the structure of nconvolution instead of maxpool can speed up the training and convergence of the network, control the gradient explosion, prevent the gradient from disappearing and prevent over fitting. At the same time, the concept structure is introduced into covseg net model, which stacks the convolution layer of 1 × 1, 3 × 3 and 5 × 5 and the pooling layer of 3 × 3 together. On the one hand, it increases the width of the network, on the other hand, it increases the adaptability of the network to scale.

We use Focal-Tversky Unet model and nn-Unet model reported in isbi as comparison models (Table 2). Through comparison, it is found that in the focus segmentation task of COVID-19 CT image, the evaluation indexes such as Dice and Precision of COVSeg-NET model proposed in this study are better than the other two models. Especially for the situation that the number of samples in medical images is small and the problem space is small, COVSeg-NET model can achieve the goal of medical image segmentation by using a variety of ways of data set amplification and preprocessing.

The trained SegNet-Lung segmentation model is used to segment all lung CT images to be classified (the remaining 29,278 lung CT images without segmentation annotations), and the segmentation marks are automatically labeled. Then use a self-editing program to extract the lung area in the image. Some examples of image annotation are shown in Figure 6.

3.   CapsNet-COVID19 classification model and lung CT image classification

Next, the CapsNet-COVID19 classification model will be used to classify the pre-segmented lung CT images. CapsNet-COVID19 has made many adjustments and improvements to the CapsNet model.

3.1.   Introduction to CapsNet model

In the CapsNet model ^[19], the capsule can contain the attributes and characteristics of the entity. The training goal of the capsule is to capture the possibility of features and their variants, not the characteristics of specific variants. Therefore, the capsule can be used not only to detect image features, but also to learn image variants. In this way, the same capsule can detect the same object category with different directions.

The input ${u}_{i}$ of the capsule is a vector, apply a transformation matrix ${W}_{ij}$ to get the output of the upper layer of the capsule. For an m × k matrix, the k-dimensional ${u}_{i}$ is converted to the m-dimensional $\hat{u}_{j \mid i} \quad((m \times k) \times(k \times 1) \Rightarrow m \times 1)$, and then the weight ${c}_{ij}$ is used to calculate the weighted sum ${s}_{ij}$:

${c}_{ij}$ is the coupling coefficient trained by iterative dynamic routing, and ${\sum }_{j}{c}_{ij}$ is designed as a sum. This model does not use the ReLU function, but uses a squashing function to scale the vector between 0 and 1:

The squashing function sets the small vector as a zero vector and the long vector as a unit vector:

Dynamic routing is located between the capsule layers, that is, between PrimaryCaps and DigitCaps in the structure diagram. This algorithm describes the calculation method of the connection weight between the $\ell$th layer and the $(\ell+1)$th layer capsule, and calculates the correlation score ${b}_{ij}$ according to the similarity:

The coupling coefficient ${c}_{ij}$ is calculated by the softmax function of ${b}_{ij}$:

After multiple iterations and updates, the obtained ${b}_{ij}$ is more accurate:

3.2.   CapsNet-COVID19 classification model research and experiment

3.2.1.   Data preparation

The 29,278 lung CT images processed by lung region segmentation can be used as the train/test data set of the CapsNet-COVID19 classification model (as shown in Figure 8 and Table 3). In order to train and test the CapsNet-COVID19 classification model, 20,000 images of the data set are used as the train set, and the remaining 9278 images are used as the test set and validationset, and the image size is set to 32 × 32.

3.2.2.   CapsNet-COVID19 classification model research and experiment

Since lung CT images are more difficult to identify and classify than ordinary images, the CapsNet-COVID19 classification model needs to be continuously modified during training to achieve better classification results. Tables 4 and 5 show the setting of the parameters of the CapsNet-COVID19 classification model.

The classification process of the CapsNet-COVID19 classification model is as follows:

(1) Input the segmented lung CT image training data set into the CapsNet-COVID19 classification model. The image set enters Conv1 first, passes through 256 9 × 9 kernel convolutions, and outputs feature maps of 256 channels (Figure 9);

(2) Then the data enters the PrimaryCaps layer. After the kernel is 5 × 5 convolution, the image size is reduced from 24 × 24 to 10 × 10. The output shape of the PrimaryCaps layer is 10 × 10 × 16 × 16;

(3) Then the data enters the DigiCaps layer, and the 16-dimensional capsule is converted into a 32-dimensional capsule using the 16 × 32 transformation matrix ${W}_{ij}$. The output of the DigiCaps layer is 3 × 32, which means 3 types of 32-dimensional capsules;

(4) Finally, the data is input to the squashing nonlinear function, and the vector ${v}_{j}$. is output. Take each vector ${v}_{j}$. as a capsule of class j, where the length of ${v}_{j}$. represents the probability of the image being recognized as each class. Finally, after the fully connected FC layer, a 32 × 32 image can be reconstructed. The reconstructed image is shown in Figure 10.

4.   Evaluation and discussion

The hardware and software configuration used in the lung CT image classification experiment is as follows:

(1) Hard disk: 16 GB RAM, 1 TB + 256 GB;

(2) Graphics processor: NVIDIA Geforce GTX 1070 8 G;

(3) Central processing unit: Intel Core i7-6700HQ processor;

(4) Operating system: Ubuntu 16.04 operating system;

(5) Software: Python, Tensorflow, Cuda8.0, NumPy, docopt, Sklearn, Matplotlib.

4.1.   Evaluation

The training of the CapsNet-COVID19 classification model on the GPU takes about 2 hours and 24 minutes. The confusion matrix can be used to evaluate the classification effect of the CapsNet-COVID19 classification model, as shown in Figure 11 and Table 6. Among them, True Positive (TP): predict the number of positive cases predicted as positive; False Negative (FN): predict the number of positive cases predicted as negative; False Positive (FP): predict the number of negative cases predicted as positive; True Negative (TN): predict the number of negative cases as negative cases. In a predictive classification model, it is usually hoped that the more accurate the predictive result, the better. That is, the larger the observed and predicted values of TP and TN, the better, on the contrary, the smaller the observed values of FP and FN, the better.

4.2.   Discussion of loss function

In addition, the evaluation indicators of the CapsNet-COVID19 classification model also include accuracy, boundary loss, reconstruction loss, and total loss. Among them, the accuracy rate refers to the ratio of the correct recognition of lung CT image categories to the total number of lung CT images:

Margin Loss: For each category $c$ that appears in the lung CT image, the capsule uses a separate loss function $L_{c}$:

Among them, if there are $T_{c} = 1, \quad m^{+} = 0.9$, and $m^{-} = 0.1$ for the object of the class, the downward weighting of λ (the default is 0.5) prevents the initial learning and shrinks the activity vectors of all classes. The length of the instantiation output vector represents the probability of the existence of each capsule entity. Only when there is a lung CT image type in the input image, this type $c$ has the longest vector output.

Reconstruction Loss: It is mainly used to promote the encoding of the input images by the DigitCaps layer:

Among them, R represents reconstruction loss, decoded images represents the encoded images, and origin images represents the original images.

Total Loss is the sum of boundary loss and reconstruction loss:

Among them, F is the total loss, λ = 0.392, therefore, the boundary loss should always dominate the reconstruction loss. If the reconstruction loss is greater than the total loss, the model will try to accurately match the output image with the input image of the train data set, resulting in an overfitting of the model and the train data.

Due to the large noise interference of the lung CT image data set used in this research, the feature difference between lung CT image categories is not obvious enough, and the image recognition is low. Therefore, when the lung CT image classification model is initially trained, the model is prone to overfitting during training.

The reasons for the over fitting of the model are as follows: (1) the amount of data in the training set is not large enough, and the training set cannot evaluate the whole data distribution, resulting in the over fitting of the model; (2) The classification model is over iterative trained in the training process, which leads to the model fitting the noise in the data and the non representative characteristics in the samples, which also leads to the over fitting of the training model; (3) The noise interference of image data set is large, the feature difference between lung CT image categories is not obvious, and the image recognition degree is low, which makes the medical image more difficult to recognize than ordinary images, resulting in over fitting of the model.

In order to avoid or reduce the over fitting phenomenon of CapsNet-COVID19 classification model, the main methods adopted in this study are as follows: first, we widely collect open data sets to ensure that the number of images in the training set is as large as possible, so as to prevent the over fitting problem. Second, in the process of training CapsNet-COVID19 classification model, when the number of iterations is set to be large and the performance of the model in the validation set/test set begins to decline, early stop measures shall be implemented to prevent over fitting caused by continuous training. At the same time, save a copy of the network model corresponding to the current number of iterations and use it as the final/optimal network model.

As shown in Figure 12, during the training process of the CapsNet-COVID19 classification model, the accuracy of the train set and the test set continue to increase as the number of iterations increases. Among them, the red polyline represents the change of the accuracy of the train set, and the blue polyline represents the change of the accuracy of the test set. As the number of iterations increased, the accuracy of the train set approached 100%. In the later stage, the accuracy of the test set fluctuated between 80% and 98%, and the highest value reached 98% accuracy at 85k step. At the same time, as the number of iterations increases, the loss value gradually decreases and tends to stabilize. At the 130 k step position, an early stop was adopted to prevent over-fitting. It should be noted that the CapsNet-COVID19 classification model gives a separate boundary loss for each capsule of the characterization category, and the total loss is the sum of the losses of all categories. The overall evaluation index of the CapsNet-COVID19 classification model is: the accuracy of the train set is 100%, the accuracy of the test set is 84.291%, and the total loss value is 11.066%. It can be seen that the CapsNet-COVID19 classification model is a robust model (Table 7).

4.3.   Comparison

We compare the performance of the CapsNet-COVID19 classification model with several better lung CT image classification methods, as shown in Table 8. The CapsNet-COVID19 classification model has a clear advantage in accuracy over the other models mentioned in the table. The number of images that the CapsNet-COVID19 classification model can process at the same time also has greater advantages than other methods. Our model can meet the requirements of massive lung CT image analysis and improve the efficiency of image analysis.

4.4.   Classification prediction

Using CapsNet-COVID19 classification model can achieve classification prediction. Input the lung CT image into the CapsNet-COVID19 classification model, and then the type and Softmax value of the image can be predicted. As shown in Figure 13. The classification method of this research is suitable for complicated medical image classification problems such as complex background, low recognition degree, blurred boundary and large image noise.

5.   Conclusions

In this study, a lung CT image classification method based on improved deep learning model is proposed. The main work and results are as follows: firstly, the SegNet-Lung segmentation model is used to segment and predict 29,278 lung CT images. Among them, SegNet-Lung segmentation model uses the structure of nconvolution (convolution2d + batchnormalization + ELU) instead of maxpool, which can speed up the training and convergence of the network, control the gradient explosion, prevent the gradient from disappearing and prevent over fitting. Secondly, the CapsNet-COVID19 classification model is used to classify and predict 29,278 segmented and preprocessed lung CT images with high accuracy. Among them, the CapsNet-COVID19 classification model improves the capsule shape of the primarycaps layer and digicaps layer in the capsule classification model. Thirdly, the evaluation method of classification model is proposed, including accuracy, boundary loss, reconstruction loss, total loss and other indicators, and the over fitting problem is solved by using early stop technology and increasing training data set. The accuracy of the test set of the classification model proposed in this study is higher than that of 169 layer densenet, inception V3, resnet50, m-inception and other methods, and once the model is trained, it can realize the rapid prediction of lung CT test images. The classification model improves the analysis efficiency of lung CT images and can meet the requirements of massive image analysis.

Conflict of interest

We declare that we do not have any commercial or associative interest that represents a conflict of interest in connection with the work submitted.