A Hybrid Capsule Network for Pneumonia Detection Using Image Augmentation Based on Generative Adversarial Network

The pneumonia dataset used in this study encompasses 5,856 chest X-ray images [30] on retrospective cohorts of pediatric patients aged 1-5 in Guangzhou Women and Children’s Medical Center. All of them were obtained as part of the patients’ routine clinical care.

The dataset contains 1,583 normal images, and 4,273 pneumonia images. The images (format: JPEG; color space: RGB) are of various sizes. The total size of the dataset amounts to 1GB.

To train the AL system, the chest X-ray images were labelled by two specialist doctors, and divided into a training set (5,216), a cross-validation set (16), and a test set (624). For quality control, all low-quality or unreadable scans were removed, such that all chest X-ray images are suitable for evaluation [31].

The capsule network, created by Sabour et al. [22], aims to prevent the loss of features in the pooling layer of convolutional neural network (CNN). The architecture of the capsule network is shown in Figure 4.

The capsule network filters the input images with a convolution layer. At the start of training, the weights of these filters are assigned randomly in the range 0-1. The convolution process determines the number and size of filters. Through the filtering, a new image can be generated by:

$x_{i, j}^{l}=\sum_{a}^{n} \sum_{b}^{n} w_{a b} y_{(i+a)(i+b)}^{l-1}$     (3)

where, y is the input image; w is the convolution filter. Spatial derivatives (e.g., horizontal, vertical, and angular edges) and color derivatives of each input image are obtained through the convolution. The convolution results are passed to the squashing function:

$v_{j}=\frac{\left\|s_{j}\right\|^{2}}{1+\left\|s_{j}\right\|^{2}} \frac{s_{j}}{\left\|s_{j}\right\|}$    (4)

where, $v_{j}$ is the output vector of the $j \text { -th }$ capsule; $s_{j}$ is the total input. This stage is called the primary capsule.

The next stage is known as the class capsule. The prediction vectors $\hat{u}_{j \mid i}$  is calculated by multiplying each input vector of this stage with the corresponding weight ( W ). Then, the total input to a capsule $S_{j}$  is calculated as the weighted sum $\hat{u}_{j \mid i}$. The output of weighted prediction vector can be expressed as:

$s_{j}=\sum_{i} c_{i j} \hat{u}_{j \mid i}, \quad \hat{u}_{j \mid i}=w_{j \mid i} u_{i}$     (5)

where, $c_{i j}$ is a coupling coefficient determined by iterative dynamic routing:

$c_{i j}=\frac{\exp \left(b_{i j}\right)}{\sum_{k} \exp \left(b_{i k}\right)}$     (6)

where, $b_{i j}$ is the log prior probability of the $i \text { -th }$ capsule to be coupled with the $j \text { -th }$ capsule. Then, an agreed scalar value is calculated between each capsule output $v_{j} \text { in layer } j$ and the predicted $\hat{u}_{i j} \text { of the } i \text { -th }$ capsule:

$a_{i j}=v_{j} \hat{u}_{j \mid i}$    (7)

The output of the class capsule layer is predicted by the dynamic routing algorithm, whose pseudocode is given in Table 2. The number of iterations of the algorithm varies with the routing constant of the dataset.

The dynamic routing algorithm calculates the lengths of class capsule outputs, which represent the class probabilities of capsule input. Then, the margin error of the capsule network is calculated by:

$\begin{aligned}

L_{k}=& T_{k} \max \left(0, m^{+}-\left\|v_{k}\right\|\right)^{2} \\

&+\Delta\left(1-T_{k}\right) \max \left(0,\left\|v_{k}\right\|-m^{-}\right)^{2}

\end{aligned}$     (8)

where, $\mathrm{T}_{\mathrm{k}}$  is 1 if each class of the image has a single value; $\mathrm{m}^{-}=0.1, \mathrm{~m}^{+}=0.9 \text { , and } \Delta=0.5$ are constants set for the pneumonia dataset. These constants depend on the size of the dataset, and the overlap conditions.

The capsule network also has a decoder module, which receives the output from each layer of the class capsule. The loss of the capsule network is defined as the sum of margin loss and decoder loss.

Table 2. Pseudocode of the dynamic routing algorithm [22]

4.png

Figure 4. Process steps in a Capsule Network Architecture [22]

5.png

Figure 5. Decoder network layers

Figure 5 is the typical structure of the decoder network. The output of the network is an image reconstructed based on the capsule values. As shown in Figure 5, the decoder network consists of class CA and FC for dealing with the chest X-ray dataset. Then, the decoder network is trained by adjusting the number of Cas and the size of mini-batch. During the training, the decoder loss is added to the CA loss multiplied by 0.005. This is because the loss of the decoder layer does not change with CA loss [22]. Then, all weights of the network are updated by the Adam optimization [36].

The experiments were carried out under the framework of TensorFlow and Keras, using the programming language of Python. In the first experiment, the CA was trained on the original dataset augmented by the classical approach. Accordingly, the chest X-ray dataset was divided into a training set and a test set. The former includes 1,583 normal images, and 4,273 pneumonia images; the latter contains 234 normal images, and 390 pneumonia images. The test dataset was determined by the data provider [31].

The CA used in the experiment was trained with constants $\mathrm{m}^{-}=0.1, \mathrm{~m}^{+}=0.9, \text { and } \Delta=0.5$. The cost of the CA architecture depends on whether the image is classified correctly. If yes, the capsule output is the minimum probability with $\mathrm{m}^{+}$ parameter; otherwise, the output is the $\mathrm{m}^{-}$ parameter. Since the probability falls between 0 and 1, $\mathrm{m}^{+} \text {and } \mathrm{m}^{-}$ were set to 0.9 and 0.1, respectively. The decay $\Delta$  coefficient provides the refutation of the cost for negative classification result. The empirical value of this coefficient is 0.5. In addition, the batch sizes (32 and 64) for training were examined, and the learning rate was selected as 0.001. The CA was trained with random weights for 50 epochs. The training was carried out 5 times, and the mean results were calculated.

The experimental results were obtained using a server computer with Intel Xeon ® 2.2 GHz processor and 64GB RAM, operating on Windows Server 2012 R2. The computer programs are run on Nvidia Quadro M4000 8GB video RAM GPU card.

The performance of our algorithm and the classical method was measured by Accuracy (Acc), F-score, Precision (Pr), and Sensitivity (Se) [37]:

$A C C=\frac{T P+T N}{(T P+F N)+(F P+T N)}$    (9)

$S e=\frac{T P}{T P+F N}$     (10)

$P r=\frac{T P}{T P+F P}$    (11)

$F-\text { score }=\frac{2 x T P}{(2 x T P+F P+F N)}$    (12)

where, TP and TF are the number of correctly classified pneumonia images and normal images, respectively; FP and FN are the incorrectly classified pneumonia images and normal images, respectively. Acc, Pr, Se and F-score metrics are derived from the confusion matrix. Acc is the ratio of the number of correctly predicted images of the total number of images. Pr and Se are the precision and sensitivity values of pneumonia disease detection, respectively. The higher these values, the better the pneumonia disease is detected. The F-score is the harmonic mean Pr and Se.

A total of five experiments were carried out. In the first experiment, the dataset was expanded with the classical data augmentation approach to examine the classification performance of the CA. The shifting (1%), rotation (random in 0-2π), and cropping (1%) operations were implemented. Then, the CA training was initiated with the batch sizes of 32 and 64. Table 3 shows the classification results of the CA with different filter sizes.

Table 3. CA classification results with different filter/kernel sizes

As shown in Table 3, the classification accuracy peaked at the filter size of 3×3. A relatively large filter size reduces the scale of the image inputted to the subsequent layers, causing a degradation of image quality. This suppresses the classification accuracy of the network.

In the second experiment, the most accurate capsule network (filter/kernel size: 3×3; batch size: 32; number of dynamic routings: 3) in the first experiment was applied with the DCGAN. The synthetic images (1,800 normal and 1,800 pneumonia) were added to the original dataset. Then, the CA was trained by the parameters of different filters and different number of convolutions. The results obtained are given in Table 4.

Table 4. Result obtained for CA with different filter sizes and increased DCGAN data

As shown in Table 4, the number of filters used affected the total accuracy. The total accuracy did not change significantly, using a single convection layer and 1,024 filters. The number of convolution layers was doubled from that of the first experiment, and retraining was carried out. In this way, the highest accuracy was obtained. Hence, DCGAN data enhancement improves the classification accuracy by 5% than the classical data augmentation approach.

In the third experiment, the parameters of the convolution layer and filter numbers were fixed to analyze how the number of images synthetized by DCGAN affects the accuracy of the CA network. The relevant results are given in Table 5.

Table 5. CA different filter number and DCGAN data increase and results obtained

As shown in Table 5, the classification accuracy of the CA increased with the number of data. After this stage, the following numbers were all increased: the number of training data, the number of convolution layers, and the number of filters. After a certain level, it was seen that the increase in the number of layers, the number of filters and the number of sample images in the training set did not affect the test accuracy. Figure 7 presents the training loss, training accuracy, and validation accuracy of the decoder layer in the CA. It can be observed that the network achieved an accuracy of 93.91%.

In the fourth experiment, the classification performance of the proposed hybrid CA-DCGAN was compared with that of two popular classifiers, namely, AlexNet [17] and GoogleNet [29]. AlexNet has 5 convolution layers and 3 FC layers, and receives images of the size 227×227×3. The first three layers of the network are responsible to extract a general feature. The weights of these layers are transferred from the trained architecture. After being filtered through these layers, the inputs are passed through two FC layers, and classified by softmax. GoogleNet is a CNN of 22 layers. The last classification layer of this network is first removed from the general structure, and substituted by a layer trained for the classification of chest X-ray dataset. As such, the final classification layer of the network is retrained. The first layers of the network are responsible for feature extraction. For this reason, the weights of the first 10 layers are kept constant in the training stage. AlexNet, GoogleNet, and CA-DCGAN were trained for 50 epochs, at the learning rate of 0.001 and the batch size of 32. The classification accuracy and confusion matrix are recorded in Table 6 and Figure 8, respectively. In the confusion matrix, the values of 1 and 2 represent the normal class and the pneumonia class, respectively. Accuracy alone cannot fully demonstrate the classification performance. Therefore, two additional metrics were adopted, including the precision and sensitivity.

7a.png

7b.png

Figure 7. The proposed DCGAN-CA classification architecture results

Table 6. CA-DCGAN, AlexNet and GoogleNET results

It can be seen that GoogleNet achieved the best FP, AlexNet realized the best FN, and our network boasted the best accuracy, sensitivity, and precision.

8.png

Figure 8. The confusion matrix obtained by the networks (CA-DCGAN, GoogleNET and AlexNET)

In the fifth experiment, our method was compared with existing methods for pneumonia classification. The classification performance is presented in Table 7. The first of the contrastive methods is CNN-based transfer learning proposed by Kermany et al. [30]. This method uses the weights transferred from the Inception3 model. After being trained for 100 epochs with a learning rate of 0.001, this method ended up with an accuracy of 92.8%. The other contrastive method is developed by Livieris et al. [38]. This method operates with the maximum likelihood selection scheme and semi-educated learning model. The proposed CA-DCGAN facilitates image classification by reducing the large images to the size of 28×28×1. As a result, our method classified low-dimensional images with a success rate of approximately 93.91%.

Table 7. Comparison with studies in the published literature on pneumonia classification