This project evaluates the effectiveness of three lightweight CNN models in detecting pneumonia in chest X-ray images, considering both performance & computational efficiency. The images are categorized into three classes: Bacterial Pneumonia, Viral Pneumonia, or Normal.
Each model employs a corresponding lightweight convolutional neural network (CNN) architecture fine-tuned through transfer learning to do multiclass image classification on a chest X-ray dataset:
- MobileNet-V2
- ShuffleNet-V2
- SqueezeNet
For a fair evaluation, all models have been trained using identical batch sizes, epochs, & data preprocessing techniques. Additionally, while the architectures of MobileNetV2, ShuffleNetV2, & SqueezeNet vary in design principles & specific layers, the process of adapting them for chest X-ray image classification remains consistent across all three models. Specifically, only the final fully connected layer of each model is replaced to align with the number of classes in the dataset. This modification enables a fair comparison of the models' performances without further alteration to their respective architectures.
The dataset is structured into three main directories: train, val, & test. Within each directory, there are subfolders representing different image categories, namely Bacterial Pneumonia, Viral Pneumonia, & Normal. Altogether, the dataset comprises 4,353 chest X-ray images in JPEG format, distributed across the three classes in each set:
- Training Set = 3559
- Bacterial Pneumonia = 1230
- Viral Pneumonia = 988
- Normal = 1341
- Validation Set = 310
- Bacterial Pneumonia = 108
- Viral Pneumonia = 104
- Normal = 98
- Test Set = 484
- Bacterial Pneumonia = 172
- Viral Pneumonia = 148
- Normal = 164
These chest X-ray images were chosen from retrospective cohorts of pediatric patients aged 1-5 years old at the Guangzhou Women & Children’s Medical Center, Guangzhou. The chest X-ray imaging was conducted as part of the routine clinical care for these patients.
Source: D. Kermany, K. Zhang, & M. Goldbaum, “Labeled Optical Coherence Tomography (OCT) & Chest X-Ray Images for Classification,” data.mendeley.com, vol. 2, Jun. 2018, doi: https://doi.org/10.17632/rscbjbr9sj.2.
Download Dataset Here:
Python PyTorch scikit-learn Pandas NumPy Matplotlib Seaborn Azure ML Studio
| Model | Accuracy (%) |
|---|---|
| MobileNetV2 | 90.08 |
| ShuffleNetV2 | 88.43 |
| SqueezeNet | 56.61 |
MobileNetV2 demonstrated the highest accuracy, indicating its strong ability to generalize & perform well on new data beyond the training & validation sets. In contrast, SqueezeNet exhibited the lowest accuracy, indicating challenges in effectively generalizing to new instances.
The training accuracy of MobileNetV2 steadily increases throughout the epochs, with validation accuracy showing fluctuations, reaching its highest point at later epochs, particularly around epochs 12 to 15. However, validation accuracy also fluctuates during the training process, notably dipping around epochs 4 to 6.
Similarly, ShuffleNetV2 exhibits a gradual increase in training accuracy over the epochs, with validation accuracy fluctuating but generally maintaining a higher level compared to MobileNetV2. The fluctuations in validation accuracy are observed throughout the training epochs, particularly notable dips around epochs 4 & 14.
In contrast, SqueezeNet starts with lower training accuracy compared to MobileNetV2 & ShuffleNetV2, but it shows improvement over time. Validation accuracy remains relatively low throughout training, gradually increasing towards the later epochs, with notable improvements around epochs 6 to 9.
Overall, MobileNetV2 & ShuffleNetV2 consistently achieve higher validation accuracies than SqueezeNet, suggesting stronger generalization abilities. While SqueezeNet improves in training accuracy, it struggles to generalize to unseen data, reflected in its lower validation accuracy. This indicates that MobileNetV2 & ShuffleNetV2 are more robust models for the task at hand.
The values along the diagonal of each matrix represent the number of correctly predicted images classified as Normal, Bacterial Pneumonia, & Viral Pneumonia, indicating the True Positives (TP) for each category.
MobileNetV2
- Bacterial Pneumonia: 167 out of 172 cases correctly predicted
- Viral Pneumonia: 124 out of 148 cases correctly predicted
- Normal: 145 out of 164 cases correctly predicted
ShuffleNetV2
- Bacterial Pneumonia: 163 out of 172 cases correctly predicted
- Viral Pneumonia: 124 out of 148 cases correctly predicted
- Normal: 141 out of 164 cases correctly predicted
SqueezeNet
- Bacterial Pneumonia: 60 out of 172 cases correctly predicted
- Viral Pneumonia: 62 out of 148 cases correctly predicted
- Normal: 152 out of 164 cases correctly predicted
MobileNetV2
| Precision | Recall | F1-Score | |
|---|---|---|---|
| Bacterial Pneumonia | 0.87 | 0.97 | 0.92 |
| Viral Pneumonia | 0.97 | 0.88 | 0.93 |
| Normal | 0.87 | 0.84 | 0.85 |
ShuffleNetV2
| Precision | Recall | F1-Score | |
|---|---|---|---|
| Bacterial Pneumonia | 0.88 | 0.95 | 0.91 |
| Viral Pneumonia | 0.97 | 0.86 | 0.91 |
| Normal | 0.82 | 0.84 | 0.83 |
SqueezeNet
| Precision | Recall | F1-Score | |
|---|---|---|---|
| Bacterial Pneumonia | 0.82 | 0.35 | 0.49 |
| Viral Pneumonia | 0.46 | 0.93 | 0.61 |
| Normal | 0.79 | 0.42 | 0.55 |
MobileNetV2 demonstrates strong precision & recall scores across all categories, indicating balanced performance in correctly identifying positive cases (true positives) & effectively avoiding false positives & false negatives. ShuffleNetV2 also exhibits respectable precision & recall scores, particularly for the bacterial & viral pneumonia categories. However, SqueezeNet's classification report reveals lower precision & recall scores, especially for bacterial pneumonia, suggesting challenges in accurately identifying positive cases for this category.
| Model | Model Size (MB) | Average Inference Speed (ms/image) |
|---|---|---|
| MobileNetV2 | 8.72 | 9.30 |
| ShuffleNetV2 | 4.95 | 9.18 |
| SqueezeNet | 2.77 | 5.79 |
The model sizes were determined from the file sizes after training, & the average inference speed was calculated by predicting 10 chest X-ray images from the test set & averaging the time taken to predict a single image in milliseconds.
As observed in the table, MobileNetV2 strikes a balance between model size & inference speed. ShuffleNetV2, while smaller in size compared to MobileNetV2, exhibits a slightly faster average inference speed. On the other hand, SqueezeNet distinguishes itself with its compact model size & relatively faster average inference speed. While MobileNetV2 & ShuffleNetV2 offer a balance between model size & speed, SqueezeNet prioritizes compactness, resulting in faster inference speeds at the expense model performance.
Access the overall evaluation here:
1. Clone this repository.
Run this command in your terminal:
git clone https://github.com/m3mentomor1/Pneumonia_Detection_with_Lightweight-CNN-Models.git
(Optional: You can also Fork this repository.)
2. Go to the repository's main directory.
Run this command in your terminal:
cd Pneumonia_Detection_with_Lightweight-CNN-Models
The models are deployed on Streamlit. Click this link to access the app.