Jlichwa/Pneumonia-Detector-Models

Pneumonia Detection - Two-Stage Classification Models

This repository contains deep learning models for a two-stage chest X-ray classification pipeline trained on pediatric data.

Stage 1: Classifies images as Normal vs Pneumonia
Stage 2: Classifies Pneumonia into Viral vs Bacterial

Both stages use ResNet-50 backbone with transfer learning and custom image preprocessing techniques.

Model Summary

• Architecture: ResNet-50 with transfer learning (ImageNet pretrained weights)
• Framework: FastAI / PyTorch
• Task: Medical image classification (binary and multi-class)
• Input: Chest X-ray images (pediatric)
• Training approach: Fine-tuning with custom preprocessing, data augmentation, and threshold calibration

Model Files

This repository contains:

• set2_stage1_bacterial_viral_detector_final - Stage 1 model (Normal-Pneumonia) Python 3.12
• set2_stage1_bacterial_viral_detector_final_310 - Stage 1 model (Normal-Pneumonia) - Python 3.10 (HuggingFace Spaces compatible)
• set2_stage2_bacterial_viral_detector_final - Stage 2 model - Python 3.12
• set2_stage2_bacterial_viral_detector_final_310 - Stage 2 model - Python 3.10 (HuggingFace Spaces compatible)
• customtransform.py - CLAHE, Colormap image transofrms applied at runtime before send to model
• image_processing.py - 2 stage model usage example

Requirements

Important: This model requires the custom transforms module included in this repository.

Download: customtransform.py from this repo

The transforms include:

• Custom contrast enhancement (entropy-based)
• Specialized augmentation for medical images
• Preprocessing pipeline optimized for chest X-rays

Make sure to download and include customtransform.py in your working directory when using these models.

Dependencies

fastai
torch
torchvision
numpy
PIL

How to Use

from fastai.vision.all import *
import customtransform  # Required - download from this repo

# Load Stage 1 model (Normal vs Pneumonia)
learn_stage1 = load_learner('set2_stage2_pneumonia_detector_final_310.pkl')

# Make prediction
img = PILImage.create('chest_xray.jpg')
pred_class, pred_idx, probs = learn_stage1.predict(img)

# If Pneumonia detected, use Stage 2 (Viral vs Bacterial)
if pred_class == 'Pneumonia':
    learn_stage2 = load_learner('set2_stage2_bacterial_viral_detector_final_310.pkl')
    subtype_class, subtype_idx, subtype_probs = learn_stage2.predict(img)
    print(f"Pneumonia subtype: {subtype_class}")

Dataset

The models are trained and evaluated on the Chest X-Ray Images (Pneumonia) dataset by Paul Mooney from Kaggle.

• Total images: 5,863 pediatric chest X-rays
• Classes: Normal, Pneumonia (Viral), Pneumonia (Bacterial)
• Splits: Train, validation, and test sets
• Dataset link: https://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia

Training Details

• Base model: ResNet-50 pretrained on ImageNet
• Preprocessing:
  • Custom contrast enhancement (entropy-based)
  • Standardization and resizing
  • Data augmentation (rotation, flip, zoom, lighting)
• Optimization:
  • Loss function: Cross-entropy
  • Optimizer: Adam with learning rate scheduling
  • Validation-based early stopping
• Threshold calibration: Stage 1 Set 2 threshold tuned from 0.50 to 0.80 on validation set to optimize precision/recall trade-off

Performance Metrics

Stage 1: Normal vs Pneumonia

Set	Stage	Accuracy	Precision (Pneumonia)	Recall (Pneumonia)	F1-score (Pneumonia)	Confusion Matrix (TN, FP / FN, TP)
Set 1	Stage 1	0.806	0.767	0.990	0.865	117, 117 / 4, 386
Set 2	Stage 1	0.848	0.804	1.000	0.891	139, 95 / 0, 390

Threshold Calibration (Set 2, Stage 1)

We tuned the decision threshold for the Pneumonia class on the validation set to improve precision while maintaining high recall.

Threshold Setting	Precision (Pneumonia)	Recall (Pneumonia)	F1-score (Pneumonia)
Before (t = 0.50)	0.625	1.000	0.769
After calibration (t = 0.80)	0.760	0.956	0.847

Raising the threshold from 0.50 to 0.80 increases precision and slightly reduces recall, resulting in a higher F1-score and fewer false positives for Pneumonia detection.

Stage 2: Viral vs Bacterial (among Pneumonia cases only)

Set	Stage	Accuracy	Macro Precision	Macro Recall	Macro F1-score	Confusion Matrix (TN, FP / FN, TP)
Set 1	Stage 2	0.897	0.926	0.866	0.884	241, 1 / 39, 109
Set 2	Stage 2	0.887	0.905	0.859	0.874	236, 6 / 38, 110

Metrics for Stage 2 are macro-averaged across the Viral and Bacterial classes. Confusion matrices show strong performance in distinguishing between pneumonia subtypes.

Intended Use

• Primary use: Research and education on medical image classification, model development, and evaluation
• NOT for clinical use: These models are not approved medical devices and must not be used for diagnosis or treatment decisions without oversight from qualified medical professionals and appropriate regulatory clearance

Ethical Considerations and Limitations

• Population bias: Dataset is pediatric-only and may not generalize to adult or geriatric populations
• Data source bias: Images from limited institutional sources may not capture full variability of real-world clinical practice
• Scanner variability: Performance may vary with different X-ray machines, imaging protocols, or clinical settings
• Risk of misuse:
  • False negatives could delay necessary treatment
  • False positives could lead to unnecessary anxiety or further testing
  • Models should only be used in controlled research settings with expert oversight forms module for preprocessing

Citation

If you use these models or this repository, please cite:

@software{lichwa_pneumonia_detector_2025,
  author  = {Lichwa, Jack},
  title   = {Pneumonia Detection - Two-Stage Classification Models},
  year    = {2025},
  version = {1.0},
  url     = {https://huggingface.co/Jlichwa/Pneumonia-Detector-Models}
}

And also cite the dataset:

@dataset{mooney2018chestxraypneumonia,
  author = {Mooney, Paul},
  title  = {Chest X-Ray Images (Pneumonia)},
  year   = {2018},
  note   = {Kaggle dataset},
  url    = {https://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia}
}

References

1. Howard, J., & Gugger, S. (2020). Deep Learning for Coders with fastai and PyTorch. O'Reilly Media.
2. Waheed, S., Ghosh, S., & Gadekallu, T. R. (2022). Pre-processing methods in chest X-ray image classification. Frontiers in Medicine, 9, 898289.
3. He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (pp. 770–778).
4. Lim, H.-W., et al. (2017). Automatic X-ray image contrast enhancement based on parameter optimization using entropy. Medical Physics, 44(5), 2212–2226.

License

MIT License - See LICENSE file for details.

Contact

For questions or collaborations, please open an issue on this repository or reach out via the Hugging Face community forum.