In recent years, large sets of radiology images have been made public. The availability of such datasets helped crowd-source the development and evaluation of deep learning models [3] [4] [5] [6]. In 2017, NIH Clinical Center released over 100,000 chest x-ray images, which comprises 108,948 frontal-view X-ray images of 32,717 unique to the scientific community. There were several studies based on this data. Wang et al. [7] through their paper “ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on weakly-Supervised Classification and Localization of Common Thorax Diseases” demonstrated that the commonly occurring thoracic diseases can be detected and even spatially located via a unified weakly supervised multi-label image classification and disease localization framework.

In 2019, MIT published MIMIC-Chest X-Ray Database (MIMIC-CXR) collection of more than 350,000 chest x-ray images associated with 227,943 studies. These images consist of both frontal and lateral views and were compiled from Beth Israel Deaconess Medical Center in Boston during years 2011 to 2016. Images are provided with 14 labels derived from free-text radiology reports using NLP tools [8].

CheXpert is another such dataset consisting of 224,316 chest radiographs of 65,240 patients who underwent a radiographic examination from Stanford University Medical Center between October 2002 and July 2017, in both inpatient and outpatient centers. Based on associated radiology reports, these X-rays images were labeled as positive, negative, or uncertain for the presence of 14 common chest radiographic observations. CheXpert data has attracted strong attention in building pre-trained learning models to address the challenges in X-Ray image processing, classification and segmentation. CheXNet is one prominent project by Stanford ML Group which is considered to have state of the art performance in classifying diseases even better than expert radiologists.

In this paper, we take inspiration from state-of-the-art CheXNet model to train similar models based on CNN that perform multi-class, multi-label classification with transfer learning from models trained on Imagenet data. Irvin et al. [1], in their paper developed models to predict five lungs’ conditions such as Atelectasis, Cardiomegaly Consolidation, Edema and Pleural Effusion. Their study achieved AUC ranging from 0.85 to 0.93. Our goal is to expand the predictions from 5 diseases to 14 diseases and achieve comparable AUCs, and also attempt to compare various pre-trained CNN models and their performance.

CNN is a deep Learning algorithm which can take an input image, assign importance as weights to various features in the image and be able to differentiate one from the other. The main advantage of Convolutional Neural network is that it has the capability to capture the temporal and spatial dependencies in an image by applying relevant filters [9]. CNN architecture is composed of a convolutional layer, a pooling layer and a fully connected layer. The function of a convolutional layer is to extract features from images. Each convolutional layer can have multiple convolution kernels and the convolutional layer is calculated as follows:

x j l = f ( ∑ i ∈ M j x j l − 1 ∗ k i j l + b j l ) ,

where x j l − 1 is the characteristic map of the output of previous layer, x j l is the output of the ith channel of the jth convolutional layer and f(.) is called the activation function. M j is the subset of input feature maps, k i j l is a convolutional kernel and b j l is its corresponding weight. A pooling layer sits between two convolutional layers to reduce the feature map dimensions and still maintain the scale of the features to an extent. Mean pooling and Max pooling are two main pooling methods used. A fully connected layer integrates multiple image maps after they are passed through multiple convolutional and pooling layers to obtain high-level semantics used to determine the class labels.

All the experiments in our work involve millions of parameters for multiple features and multiple layers of neural networks. This requires us to use a faster processing machine with graphic processing units (GPU) and cloud architecture that have pre-configured drivers and come with popular Python packages to build neural networks.

We used AWS Deep Learning AMI (Ubuntu 18.04) for training and prediction. These AMIs are pre-configured with PyTorch 1.7.1 and Python3.7 (CUDA 11.1 and Intel MKL) along with other basic data science software like pandas, numpy, scipy etc. We used g4dn.2xlarge machine with 1 Tesla T4 GPU.

We used a version of Chexpert data with down sampled resolution to train the model. The original data is 439 GB in size whereas the down sampled version is just 11 GB. We chose the down sampled version to save training time and computing power. Sample images are shown in Figure 1.

Each individual X-ray radiograph is represented by the image and a feature vector containing patient id, sex, age, study number, X-ray view-type, and a vector of fourteen expert-labeled observations. We clean the data to omit the sex and age features from the equation because we evaluate solely on the image.

Train Data Distribution

This subset of dataset is split into train and validation sets for model evaluation with a split proportion of 80:20. Figure 2 and Figure 3 denote the distribution of train data and train labels respectively. An unseen test data set is used to evaluate predictions of classes by the trained models.

Chest X-ray dataset contains images that are hard to classify for a certain disease to be present. So, radiologists often leave an uncertain label (value = −1) on such images. We followed approach similar to Irvin et al. [4] to handle uncertainties. Labels were split into u-zero and u-one categories based on previous studies. Category u-one means uncertain labels will be treated as positive and u-zero means negative. We considered labels Atelectasis, Edema as u-ones and rest of the labels as u-zero.

Dataset used for training and test are pre-processed through augmentation methods. This was done to increase the size and quality of the dataset. This process helps in solving problems related to overfitting and enhances the model’s generalization ability during training and prediction of class labels.

Based on our research, CNN architecture performs better on multi-class, multi-label classification of image dataset due to the reduction in number of parameters involved, without losing features that are critical for getting a good prediction. So, we investigated multiple models based on CNN architecture that will be discussed in detail further. Each of this model is run on train and test data with a batch size of 96 for up to 40 epochs, binary cross entropy as loss function, Adam optimizer with an initial learning rate of 0.001 which is multiplied by 10 each time the validation loss plateau after an epoch.

Accuracy: It is the ratio of number of correct predictions to the total number of input samples.

Area Under Curve: AUC of a classifier is equal to the probability that the classifier will rank a randomly chosen positive example higher than a randomly chosen negative example [14].

F1 Score: F1 Score is the Harmonic Mean between precision and recall.

Precision: The precision is the ratio the number of true positives to the sum of true positives and false positives. The precision denotes the ability of the classifier not to label as positive a sample that is negative.

Recall: The recall is the ratio the number of true positives to the sum of true positives and false negatives. The recall denotes the ability of the classifier to find all the positive samples.

We did a random 20% split of train data as validation data and trained model for a maximum of 40 epochs. ROC was used to determine early stopping criteria. Most of the models achieved best performance with 20 to 25 epochs. Figure 6 demonstrates variations in training metrics per epoch.

DenseNet model achieved highest training AUROC 78 and highest training accuracy 87% as indicated in Figure 7.

Model performance validated against success metrics. Details available are in Figure 9.

Based on overall test metrics, DenseNet121 achieved the best performance. This model achieved ROC score of 0.78 and accuracy of 87%. Other models ROC values ranged from 0.69 to 0.75 and accuracy from 83% to 86%.

Densenet121 performed better in the case of individual labels also. Figure 10 shows AUROC values for various labels. Densenet121 model achieved AUROC ranging from 0.82 to 0.93 for the 5 labels those were part of the study conducted by Stanford ML group [1]. The performance was good for some other labels also like Pleural Other (AUROC: 0.97) and Lung Opacity (AUROC: 0.91). The worst performance was for Enlarged Cardiomediastinum (AUROC: 0.49).

Densenet121 gave the best accuracy for individual labels. Figure 11 shows Accuracy values for various labels. Fracture, Lung Lesion, Pleural Other, Pneumonia and Pneumothorax achieved more than 95% accuracy. Enlarged Cardiomediastinum had the least accuracy of 53%.

Accuracy results may be misleading in some cases when the data set is unbalanced. We used confusion matrix to visualize true positives, false positives, true negatives, and false negatives as shown in Figure 12. Confusion Matrices for the best model, CustomNet and Densenet121 is given below.

Figure 13 demonstrates confusion matrix for DenseNet. The inability of the models to predict certain diseases are evident from the confusion matrix. For example, there were 68 positive cases for “Cardiomegaly” in the test data. Custom net predicted only 8 cases correctly where as DenseNet121 predicted 3 cases.