Deep convolutional neural network based medical image classification for disease diagnosis

Paredes et al. [21] use small patches of medical images as local features and k-nearest neighbor (k-NN) to classify the categorization of the whole medical image, finally achieving start-of-art accuracy. Parveen and Sathik [22] researched to detect Pneumonia from X-rays. The authors extracted features by discrete wavelet transform (DWT), wavelet frame transform (WFT) moreover, wavelet packet transform (WPT) and used Fuzzy C-means to detect Pneumonia. Caicedo et al. [23] use scale-invariant feature transform (SIFT) as a local feature descriptor and use support vector machines (SVM) classifiers to classify medical images and get state-of-art precision at 67%. However, SIFT is a patent algorithm. Thus, Rublee et al. [24] propose a free, faster local feature descriptor-oriented fast and rotated binary robust independent elementary features (ORB), which has the same performance as SIFT and even better performance than SIFT under some condition. SVM is also a high-performance classification algorithm, widely used in different medical image classification tasks by other researchers, and achieves an excellent performance [25, 26]. Therefore, this report uses ORB and SVM as the representation of the traditional methods.

With the different CNN-based deep neural networks developed and achieved a significant result on ImageNet Challenger, which is the most significant image classification and segmentation challenge in the image analyzing field [27]. The CNN-based deep neural system is widely used in the medical classification task. CNN is an excellent feature extractor, therefore utilizing it to classify medical images can avoid complicated and expensive feature engineering. Qing et al. [28] presented a customized CNN with shallow ConvLayer to classify image patches of lung disease. The authors also found that the system can be generalized to other medical image datasets. Moreover, in other research, it also found that CNN based system can be trained from big chest X-ray (CXR) film dataset and state-of-art with high accuracy and sensitivity results on their dataset, like Stanford Normal Radiology Diagnostic Dataset containing more than 400,000 CXR and a new CXR database (ChestX-ray8), which consist of 108,948 frontal-view CXR [29]. Moreover, using limited data makes it hard to train an adequate model. Therefore the transfer learning of CNN is wildly used in medical image classification tasks. Kermany et al. [3] use InceptionV3 with ImageNet trained weight and transfer learning on a medical image dataset containing 108,312 optical coherence tomography (OCT) images. They got an average accuracy of 96.6%, with a sensitivity of 97.8% and a specificity of 97.4%. The authors also compared the results with six human experts. Most of the experts got high sensitivity but low specificity, while the CNN-based system got high values on both sensitivity and specificity. Moreover, on the average weight error measure, the CNN-based system exceeds two human experts. The authors also verified their system on a small pneumonia dataset, including about five thousand images, and achieved an average accuracy of 92.8%, with a sensitivity of 93.2% and a specificity of 90.1%. This system finally may help in accelerating diagnosis and referral of patients and therefore introduce early treatment, resulting in an increased cure rate. Moreover, Vianna [30] also studied how to utilize transfer learning to build an X-ray image classification system that is the critical component of a computer-aided-diagnosis system. The authors found a fine-tuned transfer learning system with data augmentation effectively alleviate overfitting problem and yield a better result than two other models: training from scratch and a transfer learning model with only a retrained last classification layer.

As mentioned in the previous section, the CapsNet was invented in 2017 [16]. Therefore, the research about it is not as fruitful as CNN. However, there is still some research on applying them to the different datasets and varying fields due to its excellent feature—Equivariance. This means the spatial relationship of objects in an image is kept, and at the same time, the result does not impact the object’s orientation and size. Afshar et al. [18] applied CapsNet to classifying brain tumors on Magnetic Resonance Imaging (MRI) images and got 86.56% prediction accuracy with a modified CapsNet that reduces the feature maps from the original 256 to 64.

Moreover, Tomas and Robertas [31] presented a CapsNet based solution to classify four types of breast tissue biopsies from breast cancer histology images. They achieved 87% accuracy with the same high sensitivity. Jimenez-Sanchez et al. [5] evaluated the CapsNet on medical image challenges. The authors selected a CNN with three layers of ConvLayer as the baseline and compared CapsNet’s performance with LeNet and the baseline on four datasets, MNIST, Fashion-MNIST, mitosis detection (TUPAC16) and diabetic retinopathy detection (DIARETDB1), with three conditions: the partial subset of the dataset, the imbalanced subset of the dataset and data augmentation. The final result shows CapsNet performed better than the other two networks in a small, imbalanced dataset. Beşer et al. [32] implemented a sign language recognizing system by CapsNet and achieved 94.2% validation accuracy. Moreover, some researchers studied internal mechanics by varying network structures under different conditions. Xi et al. [33] studied the impact of different network structures on a complex dataset CIFAR10. The authors choose the following options:Finally, the authors found more ConvLayers and more models assembled, which have more effect on improving the final accuracy. Moreover, also they achieved the highest result with a 7-model assembled CapsNet with a more ConvLayer than the original version of Sabour’s. Furthermore, The CapsNet of Tomas and Robertas used to classify breast cancer increased the ConvLayer to five layers. On the other hand, Afshar et al. [18] also evaluated the different options of CapsNet. They fine-tuned the input size, number of feature maps, number of ConvLayers, capsule number in primary CapsLayer, dimension number in Primary Capsule, and the neuron number in reconstruction layers. The authors got the best results with a CapsNet having a \(64\times 64\) input image (original is \(28\times 28\)) and fewer feature map, which reduces to 64 from the original 256. Also, the authors found that increasing the routing iteration number beyond three will not improve the performance on the four datasets: MNIST, Fashion-MNIST, The Street View House Numbers (SVHN) dataset, and Canadian Institute for Advanced Research 10 (CIFAR10) dataset. From the previous reviews, it can be seen that the traditional method (SVM with ORB feature), CNN based transfer learning, and Capsule network can all use on the medical image dataset. Just looking at the value of accuracy on different datasets, CNN based transfer learning looks have better performance than the other two methods. However, they have not been compared to the same dataset. Therefore, this paper will compare their performance on the same dataset-the pneumonia dataset.

Moreover, there are so many different options when fine-tuning the parameter of those methods. The traditional method has so many features and classifying algorithms which can be evaluated. They cannot be iterated in this paper due to the limited time. As the baseline, the traditional method choose ORB as the feature and linear SVM as the classifier. As the data augmentation is a data preprocessing method that can apply to all three methods, it also will be evaluated on the traditional method. For CNN-based transfer learning, the layers of retrained ConvLayer, the complexity of classification layers, the dropout rate has significant effects on the final result. Therefore, they will be evaluated by this research. Based on the same research, the critical fact in capsule network: the number of the feature map, the number of the capsules, and the channels of the capsule will also be evaluated in this report.

The Dataset comes from the work of Kermnay et al. [34]. It contains two kinds of chest X-ray Images: NORMAL and PNEUMONIA, which are stored in two folders. In the PNEUMONIA folder, two types of specific PNEUMONIA can be recognized by the file name: BACTERIA and VIRUS. Table 1 describes the composition of the dataset. The training dataset contains 5232 X-ray images, while the testing dataset contains 624 images. In the training dataset, the image in the NORMAL class only occupies one-fourth of all data. In the testing dataset, the PNEUMONIA consists of 62.5% of all data, which means the accuracy of the testing data should higher 62.5%.

Figure 1 shows examples of chest X-rays from the dataset. The normal chest X-ray (left panel) depicts clear lungs without any areas of abnormal opacification in the image.

Bacterial pneumonia (middle) typically exhibits a focal lobar consolidation, in the right upper lobe (red rectangle), whereas viral pneumonia (right) manifests with a more diffuse interstitial pattern in both lungs.

In this paper, three data augmentation algorithms evaluated. It can be seen from Table 2, Aug0 means using the original dataset without augmentation. Aug1 means simple geometrical transform of the image: such as randomly flip horizontally and vertically, randomly rotates within 0.05\(^\circ\), horizontal shear within the range 0.05 times the image width and zoom in within 0.05 times while Aug2 is a more complicated transform than Aug1. Besides all transforms of Aug1, it also does a slightly horizontal and veridical shift. To avoid the data exploration of the combination of the data augmentation models and classification algorithms, this paper only evaluates the effects of different augmentation algorithms on VGG16. This is the best classification algorithm for this paper. However, to analyze the effects of data augmentation, all three classification algorithms also evaluated on Aug0, and the best augmentation model got from this test.

The ORB, VLAD, and SVM classification are chosen as the baseline. Two experiments conducted: First is classifying Normal and Pneumonia with the original dataset. Second does the same classification but with the best augmentation models.

Because the chest X-ray dataset is small and different from ImageNet, whose weight used in the transfer learning experiments, therefore three group experiments conducted to fine-tune the final model. The first group of experiments aims to evaluate the effect of classification layer size on the final classification accuracy. Five models used on two CNN: VGG16 and InceptionV3 showed in Table 3. In the 2nd column, the classification model described. Example, model3 consists of eight layers after the ConvLayers which are: global average pooling (GAP) layer, fully connected (FC) layer with 512 neurons, dropout layer with 50% drop rate, second FC layer with 256 neurons, second dropout layer with 50% drop rate, third FC layer with 128 neurons, third dropout layer with 50% drop rate and a classification layer with a SoftMax activation function. The last two columns list the parameters needed to be trained in VGG16 and InceptionV3. They used to indicate the complexity of the model.

The second group experiment aims to evaluate how many ConvLayers should be unfrozen and trained. A total of three experiments conducted (showed in Table 4). The first experiment evaluates the results of the best classification model with an unfrozen ConvLayer. Because the training parameter of the last ConvLayer is quite large, to prevent overfitting, the second experiment uses a smaller classification model. For testing the limit of the number of unfrozen ConvLayers, the third experiment unfreezes one more ConvLayers of the better ones in the previous two models.

The third group experiment fine-tunes other parameters based on the best model in the previous two experiment groups, such as increasing the drop rate of the dropout layer, reducing the learning rate, adding a batch normalization layer, which makes learning more stable and quicker.

For CapsNet, the feature map number, the size of the PrimaryCaps layer and input image size impact on the performance of the classification. Thus, Table 5 shows the experiments aimed at evaluating the effects of those parameters.

In Table 7, the first column is the augmentation methods, and the second column is the average accuracy of the linear SVM classifier with ORB features. It can be seen the augmentation with more images increase the accuracy.

In Table 8, it can be seen that VGG16 is better than InceptionV3 and Model3 in VGG16 is the best classification model. Thus, in the following experiment, VGG16 and Model3 continue to be fine-tuned.

From Table 9, it can be seen that the classification model2 with the last unfrozen ConvLayer was the best model in all three experiments. Thus, the following experiments will continue to be fine-tuned.

The experiments in Table10 are explorational testing. The most successful models in previous experiments: Model3 and Model2 with the last unfrozen ConvLayer chosen as the baseline. Then according to the results and the effects of dropout, learning rate, and more training data, the parameters adjusted. In all the experiments, Model2 with the last unfrozen ConvLayer and all other default parameters still has the best results (showed in Tables 9, 10).

Table 11 shows the experiments processes and results by capsule network. The augmentation first is evaluated. The Aug1 is better than no augmentation, and the number of feature maps, input size, number of primary capsules, number of capsule channels, and number of training images vary. The best result comes from tests no. 11 and 13. They both have fewer feature maps, primary capsules, and capsule channels. The difference between them is the number of feature maps. One has 24 feature maps, while the other has 32 feature maps.

To check if the findings can be used on other datasets, some experiments conducted on the OCT dataset which was published together with the Pneumonia dataset but included 108,309 OCT images. From Table 12, it can be seen, the best result obtained from test 5.

Table 13 shows a summary of all the experiment test result between no augmentation and augmentation.

It can be seen that augmentation improves performance regardless of the model. That is because augmentation geometrically transforms the picture, which facilitates the machine learning algorithm to learn the underground feature without the impact of rotation and scale. However, from Table 6, it can be seen that complicated transforms are not always better than simple ones. Too complicated transforms introduce some noise in the feature that disturbs the learning process.

To evaluate the performance of the models in this paper, Table 13 compares the best results of different models on the same pneumonia dataset. From Table 16, it can be seen that a neural network-based method is significantly better than the traditional method because it is a useful feature learner during the traditional method, just a feature-ORB. In version 2 of the dataset, the best model, VGG16 in this paper, got slightly lower accuracy and recall than the state-of-art result but obtained a higher specificity. On the latest dataset, the performance of VGG16 was generally higher. The VGG16 model released the last ConvLayer so that it would learn the specific features of the dataset. That should significantly help to improve performance very much. Kermany’s work also retrains the ConvLayer of InceptionV3, but the model overfits too much to get an excellent test performance. The reason why our model does not overfit too much maybe because the VGG16 model is not as complicated as InceptionV3.

From Table 12, it can be seen that the best model comes from test 5 instead of test 1. The new model adds complicated FC layers; therefore, the full complicity is better matched with the new dataset. The unfrozen two ConvLayers will make the system too complicated for the new dataset and, therefore, cannot find the local maxima. The best result is slightly lower than the start-of-art result of Kermany’s work (96.6%). However, this experiment result also confirms our findings. The specific feature is most important to improve accuracy—the proper model complexity help to find the best result.