Detection of single and dual pulmonary diseases using an optimized vision transformer

1.1 Problem statement

Manual methods for using CXRs to detect and distinguish between pulmonary diseases can be challenging and time-consuming. Furthermore, low-resource nations lack skilled radiologists, especially in remote areas. Several scholars have used CNN-supported computer vision to detect and diagnose pneumonia, tuberculosis, COVID-19, and lung cancer without the intervention of a radiology expert. However, in some studies, the focus has been on a narrow range of diseases, and physicians’ needs have not been considered, which restricts the clinical relevance of this literature. In addition, some researchers have relied on small datasets, and they have not paid attention to identifying more than one disease in a single image, which reduces model robustness and generalizability. The challenge is to develop a model that can predict and discriminate pulmonary diseases with high accuracy, given that some of these diseases have similar footprints in CXRs, which are difficult to distinguish for human experts (an instance of the interclass similarity problem). ViTs are the most promising tools to solve this problem but they require vast datasets to produce good results. This is one of the obstacles in classifying single- and dual-disease CXR images using a ViT.

1.2 Main contributions of this study

The main contributions of this study are as follows:

• The selection, adaptation, and optimization of a model that can help radiology experts detect and classify the most common pulmonary diseases from CXR images without the need for further examination, such as with CT or MRI.

• A system that will classify eight pulmonary diseases based on studies, including combination classes involving effusion, pneumonia, tuberculosis, pneumothorax, COVID-19, and normal status, as diagnosed by CXR.

• The use of a large balanced dataset of 70,228 frontal view CXR images. This dataset has been collected from a global source and combined with a local source from a Saudi Hospital.

• Conducting several experiments to compare state-of-the-art models, such as well-known CNN pre-trained models, ViTs, and combined pre-trained CNN models with vision transforms, in the development of an optimal model that can predict and discriminate between pulmonary diseases with high accuracy.

• An optimized model based on ViTs that outperforms state-of-the-art models.

• The use of advanced data pre-processing techniques to overcome imbalanced datasets, multi-label issues, and low contrast has been carefully considered and demonstrated to positively impact accuracy.

1.3 Article outline

The rest of this article is organized as follows: The relevant literature on the image classification techniques based on CNN and ViTs is presented in Section ‎2. Section 3 discusses disease selection, dataset preparation, model comparison conditions, and the proposed model architecture. Section ‎4 presents the results and their evaluation metrics and a comparison with other state-of-the-art models. The study’s limitations are given in Section 5. The conclusions are offered in Section ‎6.

2.1 State-of-the-art ViTs for image classification

The use of ViTs with powerful self-attention mechanisms in medical imaging-based disease diagnosis has seen significant development recently. Jiang and Lin proposed a method that combines Swin transformer and transformer in transformer for classifying individuals as having normal conditions, pneumonia, or COVID-19 based on CXR images. Their model was trained using a dataset of 22,585 and achieved a notable accuracy of 94.75%, a sensitivity of 94.75%, and a specificity of 95.09% [25]. Meanwhile, Duong et al. developed a hybrid model based on EfficientNet and ViT to address concerns about the use of heterogeneous datasets when classifying individuals as having tuberculosis, pneumonia, or normal conditions. Their hybrid model demonstrated an accuracy of 97.72%, indicating a significant improvement in disease diagnosis [26]. To enhance lung disease diagnosis, Ko et al. developed a model using various ViT models (i.e., ViT, FastViT, and CrossViT) combined with various optimizers. The combination of ViT and the RAdam optimizer, forming a model that had a balanced dataset, achieved the highest accuracy of 95.87% [27]. Angara et al. also presented a hybrid framework based on ResNet34 and ViT to detect normal conditions and pneumonia based on CXR images. This hybrid framework demonstrated a notable accuracy of 94.87% [28]. Some state-of-the-art studies on pulmonary disease classification accuracies are presented in Table 1. Table 2 presents the comparison of these works considering the percentage of the dataset size used in the training, validation, and testing phases.

Although significant studies on the classification of pulmonary diseases have been published, there are still numerous critical gaps in existing research. First, one of the limitations discovered during this literature review was that most scholars have not looked for a solution to the medical problem of classifying pulmonary diseases based on the relationship between the disease and the type of radiology used in the diagnosis process. To develop a complete system for the classification of such diseases based on CXRs, only diseases that can be diagnosed using solely CXRs, without the need for CT or MRI scans, should be studied. For instance, a CT image should be considered for the diagnosis of lung cancer or asthma, whereas COVID-19, pneumonia, pneumothorax, tuberculosis, and pleural effusion can be fully diagnosed from only a CXR. Second, some of the reviewed studies did not give prominence to image quality improvements by using techniques, such as the contrast-limited adaptive histogram equalization (CLAHE) technique, before applying the deep learning model. Third, a few scholars have discussed the possibility of multilabel classification in pulmonary imaging, where more than one disease is present in a single CXR image.

Fourth, the training datasets used in previous state-of-the-art research are small. Also, some researchers examined vast training datasets but provided very limited testing datasets. Using data that have not been previously employed and do not exhibit systems scalability can lead to partial results. Model performance and accuracy are significantly impacted by the size of the dataset. To improve the model’s generalization, robustness, and efficacy, it is important to have a large dataset to let the model learn well. This study addressed these gaps by adapting and optimizing the pre-trained ViT-B16 model to categorize eight classes of evidence concerning pulmonary diseases that can be identified using CXRs, which are inexpensive and can be used anywhere, even in remote areas. To overcome the limitations of the dataset, global and local datasets were employed, and data preprocessing techniques were utilized to solve the imbalanced dataset problem. By taking these measures, the proposed model outperformed other state-of-the-art models in terms of accuracy and generalizability, as will be demonstrated below.

3.1 Disease selection based only on CXR images

The study’s main aim was to propose a model that could detect pulmonary diseases based on CXRs without the need for CT images, which are more expensive and less available in developing countries and remote areas. Even in the absence of experts, this would help novice physicians achieve early detection and distinguish between pulmonary illnesses with high accuracy. After several meetings with the physician in the ICU of the National Guard Hospital and based on several existing studies, we selected the most common diseases seen in CXRs. One of the studies we consulted, an article by Mishra et al., described some cases of patients diagnosed as positive for COVID-19 with evidence of pneumonia [29]. The study by Xu et al. investigated a case of COVID-19 and pneumothorax observed in a CXR [30]. Grafakou et al. examined pneumonia and pleural effusion [31] in the CXRs of children. Another study, discussed in Imaging in Pediatric Pulmonology, highlights the possibility of observing both pleural effusion and pneumothorax in children [32]. To solve the problem of two diseases existing in the same image, we dealt with each combination as a new single class; however, this increased the number of classes and, hence, the difficulty of the inter-class similarity problem. According to the above literature, the proposed model classified classes of evidence as follows:

1. Class 1: Pleural effusion

2. Class 2: Pneumonia

3. Class 3: Pneumothorax

4. Class 4: Tuberculosis

5. Class 5: COVID-19

6. Class 6: Pleural effusion and pneumonia

7. Class 7: Pneumothorax and pleural effusion

8. Class 8: Normal status

3.2 Dataset collection from different sources

The dataset used in this study was collected from a global source: the National Institutes of Health CXR Dataset (ChestX-ray8), Stanford University Medical Center Shared Dataset, and Kaggle and combined with the local dataset from Saudi Hospital. This dataset included 51,889 frontal CXR images classified into eight categories: COVID-19, pneumonia, pneumothorax, tuberculosis, pleural effusion, pleural effusion and pneumonia, pneumothorax and pleural effusion, and normal. Table 3 shows the size, source, and description of each disease used in the dataset. The dataset was split randomly into 75% training, 5% validation, and 20% testing, as presented in Table 4.

3.3 Solving the problem of imbalanced datasets

One of the greatest challenges in developing and optimizing a ViT model to detect pulmonary diseases is finding a balanced large dataset that includes both CXR images with single diseases (e.g., tuberculosis) and images with two conditions (e.g., pneumothorax and pleural effusion). To overcome the limitations of imbalanced datasets, avoid overfitting, and improve model performance, we applied the data augmentation technique (i.e., the resampling technique) to minority classes, as shown in Figure 1.

Figure 1

Resampling proposed technique (created by the authors).

Before applying data augmentation to the minority classes, it was necessary to define the target number of images that had to be reached in such classes; this was the average of the number of images in all the classes, which is defined as follows:

where N is the total number of classes, and Class Count _ⅈ is the number of images in each Class (i).

After obtaining the target count, we obtained the number of images needed for each minority Class (i), as follows:

Table 5 shows the total number of images for each class before and after the data balancing process.

3.4 Expert-guided data cleaning

Because of uncertainties concerning certain diseases in some data sources, we reviewed, cleaned, and corrected the dataset with radiologists from the National Guard Hospital and a consultant physician from the Pulmonary Medicine Department to ensure accuracy and remove any incorrect data or unclear images.

3.5 Improving image quality using CLAHE

The CXR image dataset was collected from various sources. To achieve a high prediction accuracy for the classification, the images had to be clean, balanced, and enhanced. We utilized the CLAHE technique in the data preprocessing phase to enhance the images and improve their contrast [33]. CLAHE is an enhancement technique based on adaptive histogram equalization that can improve the visual contrast by expanding the intensity range or stretching out the most common intensity value found in the image [34]. The CLAHE technique splits the image into sub-images called “blocks” and then equalizes the histogram of each block separately, which helps to improve the contrast and the results in a very visible way. Figure 2 shows an example of the enhancement gained when using the CLAHE technique.

Figure 2

CXR images before and after applying the CLAHE enhancement technique (created by the authors).

Figure 3 displays the chart of the histogram of pixel distribution versus the gray levels in an image after applying the CLAHE technique. The chart shown indicates that the image has good contrast and high visibility because the pixel intensities are approximately evenly distributed across the range from 0 to 255.

Figure 3

Intensity distribution histogram after applying the CLAHE enhancement technique (created by the authors).

3.6 Improving the model’s generalization via data augmentation

To make the model more robust and improve its generalization, we utilized the data augmentation training options to increase the size of the training dataset. Table 6 shows the parameters we used for data augmentation during the training of our model.

3.7 Minimizing data imbalance with the class weight technique

In the model, class weights were calculated before training using the compute_class_weight function from the sklearn.utils.class_weight module to address class imbalance in the dataset. The formula used is:

Where:

1. N is the total number of samples in the dataset,

2. is the number of samples in class_i,

3. C is the total number of classes.

Figure 4

Computed weight distribution for each class (created by the authors).

These calculated weights were then incorporated into the training phase through the fit function. The class_weight parameter adjusts the loss function dynamically to impose higher penalties for errors involving underrepresented classes.

This modification penalizes misclassifications of minority classes more heavily, thereby encouraging the model to learn from these underrepresented instances.

3.8 Experimental setup of the model comparison conditions

To prove its performance and robustness, the proposed model was compared with the well-known CNN pre-trained models (VGG16, VGG19, ResNet50, InceptionV3, and DenseNet121). The same data were used, and all the abovementioned techniques were applied, including balancing the dataset, applying class weights to avoid bias, utilizing CLAHE, and performing data augmentation. During the comparison, it was observed that applying the same hyperparameters used for the Vision Transformer (ViT-B16) model (e.g., a learning rate of 10⁻⁶) led to difficulties for the CNN models in effectively adapting to the dataset. To address this challenge, the hyperparameters for the CNN models were adjusted. Specifically, the learning rate was increased to 10⁻³, the batch size was set to 120, and the training epochs were extended to 100. These adjustments provided the CNN models with sufficient time and capacity to adapt to the large and diverse dataset.

3.9 Customizing the proposed architecture

In these experiments, the powerful pre-trained ViT was utilized, and fine-tuning was performed to improve the accuracy of the classification of pulmonary diseases. The hyperparameters and model architecture were meticulously fine-tuned for optimal results. Nearly 50 experiments were conducted to adjust the model.

3.9.2 Experimental setup for single- and dual-disease classification

Next, two additional classes were created for dual diseases: Class 7 (pleural effusion and pneumonia) and Class 8 (pneumothorax and pleural effusion). While the proposed methodology achieved high performance in classifying single diseases, the same parameters, when applied to dual-disease classification, resulted in a reduction in accuracy to approximately 93%. This highlights the added complexity of distinguishing overlapping features in multi-disease scenarios.

To improve the accuracy of dual-disease classification, we conducted a series of experiments totaling nearly 13 experiments. The following paragraphs summarize the specific modifications and methodologies employed to address the complexities of this task.

First, we started to fine-tune the learning rate using equation (4) to obtain the best initial learning rate for this model:

(4) LR ( t ) = a × ( b / a ) t ,

where a is the initial learning rate, b is the final learning rate, and t is the time step. We concluded that for single-disease and dual-disease classification, the learning rate of 10⁻⁵ was optimal. Table 7 shows the fine-tuned hyperparameters used to optimize the output model of the training process.

Table 7

Fine-tuned hyperparameters of the proposed model

Hyperparameter	Value
Image size	224
Number of channels	3
Batch size	16
Learning rate	1 × 10−6
Number of epochs	100
Number of classes	8
Class name	Pleural effusion
	Pneumonia
	Pneumothorax
	Tuberculosis
	COVID-19
	Pleural effusion and pneumonia
	Pneumothorax and pleural effusion
	Normal status

The data balancing techniques and the fine-tuned parameters for the data augmentation were the same as explained in the preceding section. Regarding the training parameters of the ViT model, initially, to retain the low-level feature weights learnt from the huge ImageNet dataset, we froze the first five encode layers, which are the input layer, embedding layer, reshape layer, class token, and position embedding layer; this means that these layer weights will not be updated during the training, as illustrated in Figure 5.

Figure 5

Initial phase of the proposed model architecture (adapted from the study of Dosovitskiy et al. [12]).

Since a large number of datasets were used to train the proposed model, it was important to unfreeze the remaining 12 transformer encoder blocks to allow the model to learn the high-level relationships among the patches and the new image classes. This resulted in an increase in classification accuracy through the multi-head self-attention mechanism and the feedforward multi-layer perceptron function in each encoder block. This contributed to the overall enhancements of the different elements interacting in the training of the proposed model shown in Figure 6.

Figure 6

Middle and top phases of the proposed model architecture (adapted from the study of Dosovitskiy et al. [12]).

Next, a fully connected dense layer with 512 units and a ReLU activation function were added. Also, to reduce overfitting, a dropout layer with a rate of 50% was inserted. Afterward, another dense layer of 256 units was added; an activation layer was introduced, and another dropout layer of 50% was inserted. At the top of the model, the following components were added: a flattened layer, another dropout layer (with a 50% dropout rate), a final dense layer with a number of output neurons equal to the number of classes (eight), and a SoftMax layer to generate relative priorities for multiclassification. Figure 7 shows the modified ViT-B16 architecture.

Figure 7

The complete architecture of the proposed model (adapted from the study of Dosovitskiy et al. [12]).

To improve the efficiency of the training and the model’s performance, callback and learning rate scheduling were used. During the process of training, callback helps to save the model at a specific checkpoint to obtain the best version of the model and avoid any loss or deterioration because of overfitting. Additionally, EarlyStopping was employed to monitor the loss and stop the training if there was no improvement after a specific number of epochs. This helped to reduce overfitting. Finally, the ReduceLROnPlateau parameter was utilized to monitor the loss. If there was no improvement after three epochs, the learning rate was reduced by multiplying it by 0.2 (Table 8).

Table 8

Callbacks for the proposed model

Model checkpoint	Monitor: Accuracy and save best only
Early stopping	Loss
Patience	5
Reduce LR on the plateau	Monitor: Loss
Factor	0.2
Patience	3
Min_lr	1 × 10−7
Verbose	1

4.1 Results

This section presents the results of each experiment and shows how our model accurately classified various pulmonary diseases. We compared our model with the well-known CNN models and state-of-the-art models. We describe below the final experiments and present their purpose and relevant findings.

4.2 ROC–AUC scores of the models for all disease classes

The receiver operating characteristic (ROC) curve and its AUC score indicate how well a model can produce relative scores to discriminate between positive and negative cases across classification thresholds. The ROC–AUC scores range from 0 to 1, with 0.5 indicating random guessing and 1 indicating the perfect performance.

To assess the performance of the proposed ViT-B16 model compared to pre-trained CNNs, the classification of eight pulmonary disease classes was analyzed. The models’ ability to differentiate each disease class against all other conditions was evaluated. Table 9 summarizes the ROC–AUC scores for each class and the overall performance of each model. Figures 8–13 provide the ROC curves for visual comparison.

Figure 8

ROC–AUC for the VGG16 model (created by the authors).

Figure 9

ROC–AUC for the VGG19 model (created by the authors).

Figure 10

ROC–AUC for the InceptionV3 (created by the authors).

Figure 11

ROC–AUC for the DenseNet121 (created by the authors).

Figure 12

ROC–AUC for the ResNet50 (created by the authors).

Figure 13

ROC–AUC for the proposed model architecture (created by the authors).

The proposed ViT-B16 model achieved a perfect overall AUC score of 1.00, demonstrating exceptional classification accuracy across all pulmonary disease classes and outperforming pre-trained CNN models such as VGG16, InceptionV3, and DenseNet121, whose overall AUC scores ranged from 0.94 to 0.96. This superior performance can be attributed to the advanced feature extraction capabilities of the ViT, which effectively handles complex patterns in medical imaging data. Figures 8–13 illustrate the performance differences, with the ROC curve for the ViT-B16 model (Figure 13) demonstrating its ability to achieve perfect classification for all classes. These results highlight the efficacy of ViTs as a robust solution for medical imaging tasks, addressing the limitations of traditional CNN models and paving the way for further exploration and clinical deployment.

4.3 Comparison of the performances of the models

The performances of the proposed model for multi-class classification were compared against those of the pre-trained CNN models (VGG16, VGG19, InceptionV3, DenseNet121, and ResNet50). Six experiments were conducted, and the results are presented in Table 10.

Table 10 shows that the proposed ViT-B16 model significantly outperformed all other models across all metrics, achieving an impressive accuracy of 98.18%, precision of 98.11%, recall of 98.10%, and F1-score of 98.18%. The results show that the proposed model strikes a balance between recall and precision. Among the pre-trained CNN models, VGG16 showed the highest performance, achieving an accuracy of 87.78%, followed by VGG19 with 86.47%. Models such as InceptionV3, DenseNet121, and ResNet50 struggled to achieve comparable performance, with their accuracy ranging between 68.65 and 71.61%. These models exhibited lower recall and F1-scores, reflecting their challenges in balancing false positives and false negatives in multi-class pulmonary disease classification, as shown in Figures 14–18. The results highlight the superior feature extraction and classification capabilities of the ViT-B16 model, which effectively handled the complexity of multi-class classification tasks.

Figure 14

The confusion matrix illustrates the classification results on the test dataset using the VGG16 model. Each cell represents the number of samples classified into the corresponding predicted category (created by the authors).

Figure 15

The confusion matrix illustrates the classification results on the test dataset using the VGG19 model. Each cell shows the number of samples predicted for each class (created by the authors).

Figure 16

The confusion matrix illustrates the classification results on the test dataset using the InceptionV3 model. Each cell shows the number of correctly and incorrectly classified instances for each class (created by the authors).

Figure 17

The confusion matrix presents the classification results on the test dataset using the DenseNet121 model. Each cell shows the number of true versus predicted labels for each class (created by the authors).

Figure 18

The confusion matrix displays the classification outcomes for the test dataset using the ResNet50 model. Each value represents the number of samples predicted for each class (created by the authors).

In the medical field, minimizing false negatives is important, especially in serious cases, as predictions affect a patient’s life and might lead to erroneous diagnoses. The confusion matrix in Figure 19 demonstrates how our model minimizes false negatives while maximizing the true positives, which is its strength. Moreover, our model demonstrated a high recall (sensitivity) even in cases where pleural effusion co-existed with either pneumonia or pneumothorax, indicating its ability for accurate disease detection.

Figure 19

The confusion matrix summarizes the classification performance of the proposed model on the test dataset. Each cell represents the number of instances predicted for each class (created by the authors).

4.4 Class-wise performance evaluation

The ability of the proposed model to correctly classify each class in a multi-class classification was examined. Its performance was evaluated in relation to the output confusion matrix, which was constructed based on the results achieved by the model in classifying the test data.

As shown in Table 11, the highest scores for precision (i.e., the highest percentage of correct labels) were noted in pneumonia, pneumothorax, tuberculosis, and COVID-19. By contrast, a low precision score of 95% was observed when pleural effusion and pneumonia co-existed. As for recall, the highest score was noted in normal cases, indicating the proposed model’s sensitivity in identifying normal conditions. The second highest recall was 99% for tuberculosis, followed by 95% for cases in which pneumothorax and pleural effusion co-existed. The proposed model achieved a good balance between precision and recall for normal cases, pneumonia, tuberculosis, and COVID-19, achieving a high F1-score of 99%.

4.5 Comparison with other state-of-the-art models

In this subsection, we compare the performance of the proposed model against the recent state-of-the-art models based on the previous results achieved by our model and the reported results for each SOTA case. To ensure a comprehensive and fair comparison, we excluded the models with a dataset of less than 10,000 images. This threshold would ensure that all the other models have sufficient datasets for training and that the evaluation is reliable. Table 12 presents the pulmonary disease classification results of the proposed model along with the different SOTA models involving different techniques.

As shown in Table 12, our model demonstrated a greater classification accuracy than the other models, although it analyzed the highest number of classes. To the best of our knowledge, this study is the first to distinguish between eight classes, including normal conditions and seven pulmonary diseases, which occur either alone or in combination with another disease (i.e., COVID-19, pneumonia, pneumothorax, tuberculosis, pleural effusion, pneumonia, and pleural effusion, and pneumothorax and pleural effusion). Our model demonstrated an accuracy of 98.18%, and it outperformed even those that classified fewer classes and those that included fewer datasets. The models using various ViT architectures [25,26,27] achieved high accuracy, indicating the significant improvement of ViTs in medical imaging-based disease diagnosis. However, our tuned ViT model demonstrated higher accuracies than the shown ViT-based techniques, although distinguishing between more classes due to the well-suited data pre-processing techniques and the perfect tuning of our model’s architecture and its hyperparameters.