Pneumonia Chest Xray Classifier

Model Details

Model Description

The models take a single chest X-ray (anterior–posterior view) and output a class probability. Input images are resized to 224×224, converted to 3 channels, and normalised with ImageNet statistics. Transfer-learning backbones are trained in two phases (head-only, then full fine-tuning with early stopping); each model is trained with three seeds (0, 1, 2) for reproducibility.

• Developed by: Luana Carolina Reis and Jakub Błaszczyk (University of Aveiro, Portugal; Lodz University of Technology, Poland)
• Funded by: Not funded (academic course project)
• Shared by: luanacarolina
• Model type: Image classification (CNN); binary and 3-class chest-X-ray pneumonia detection
• Language(s) (NLP): N/A (image model)
• License: CC BY 4.0 (weights); dataset under CC BY 4.0
• Finetuned from model: torchvision ResNet18 and DenseNet121, ImageNet-1k pretrained weights (the custom CNN is trained from scratch)

Model Sources

• Repository: https://huggingface.co/luanacarolina/pneumonia-chest-xray-classifier (weights); project code on GitHub
• Paper: Based on Kermany et al. (2018), Identifying Medical Diagnoses and Treatable Diseases by Image-Based Deep Learning, Cell 172(5):1122–1131. https://doi.org/10.1016/j.cell.2018.02.010
• Demo: N/A

Uses

Direct Use

Research and educational use: classifying pediatric chest X-rays as NORMAL vs PNEUMONIA (binary) or NORMAL/BACTERIA/VIRUS (three-class), and reproducing the reported results. The accompanying Grad-CAM tooling can be used to visualise which lung regions drive each prediction.

Downstream Use

The checkpoints can serve as a starting point for further fine-tuning on other (e.g. adult or multi-institution) chest X-ray datasets, or as a baseline in pneumonia-detection benchmarks.

Out-of-Scope Use

Not for clinical use. This is a course project, not a medical device, and must not be used for real diagnosis, triage, or treatment decisions. It is trained on single-institution pediatric data and is not validated for adults, other scanners/populations, or imaging modalities other than AP chest X-ray. The three-class viral/bacterial distinction in particular is unreliable (see Limitations).

Bias, Risks, and Limitations

• Single-institution, pediatric data (Guangzhou Women and Children's Medical Center, ages 1–5). Generalisation to adults or other hospitals/scanners is untested.
• Class imbalance: the training set has ~2× more bacterial than viral or normal images; we mitigate with inverse-frequency class weights, but it still biases behaviour.
• Viral pneumonia is the bottleneck: its diffuse interstitial pattern overlaps with both normal lungs and early bacterial infiltrates, so VIRUS has the lowest F1 (~0.72) and drives the dominant NORMAL→VIRUS confusion. Neither class weighting, ensembling, nor a hierarchical design closed this gap.
• Backbone interpretability differs: with similar metrics, DenseNet121's Grad-CAM heatmaps are anatomically focused, whereas ResNet18 often attends to non-parenchymal structures ("right answer, wrong reason"), a generalisation risk.
• Modest specificity on the binary task (~0.62 for the ensemble): high pneumonia recall comes at the cost of false alarms on normal lungs.

Recommendations

Users (both direct and downstream) should treat outputs as research signals only, prefer DenseNet121 when an auditable explanation is needed and the ensemble when raw predictive performance matters, and validate on their own data before any deployment. Predictions for the VIRUS class should be treated with particular caution.

How to Get Started with the Model

Use the code below to get started with the model.

import torch
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
from PIL import Image

# Download the primary binary checkpoint
ckpt = hf_hub_download(
    repo_id="luanacarolina/pneumonia-chest-xray-classifier",
    filename="checkpoints/binary/densenet121.pt",
)

model = models.densenet121(weights=None)
model.classifier = torch.nn.Linear(model.classifier.in_features, 1)  # binary BCE head
state = torch.load(ckpt, map_location="cpu", weights_only=True)
model.load_state_dict(state.get("model_state_dict", state))
model.eval()

tf = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.Grayscale(num_output_channels=3),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

img = tf(Image.open("xray.jpeg").convert("RGB")).unsqueeze(0)
with torch.no_grad():
    p_pneumonia = torch.sigmoid(model(img)).item()
print(f"P(pneumonia) = {p_pneumonia:.3f}")

For the three-class models use num_classes=3 with a softmax head and the checkpoints under checkpoints/three_class/.

Training Details

Training Data

Kermany Chest X-Ray Images (Pneumonia) dataset, distributed on Kaggle by Paul Mooney: https://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia (≈5,856 AP pediatric chest X-rays; NORMAL / PNEUMONIA folders, with bacterial/viral encoded in the pneumonia filenames). Original data from Kermany et al. (2018).

Training Procedure

Preprocessing

Resize to 224×224, grayscale→3 channels, ImageNet normalisation. Training augmentation: random horizontal flip and ±10° rotation. For the three-class task we use a patient-aware stratified validation split (whole patients held out, ~12% per class) to avoid the data leakage caused by multiple X-rays per patient and by the tiny 16-image official val/ folder.

Training Hyperparameters

• Training regime: fp32
• Optimizer: Adam, learning rate 1e-4
• Batch size: 32
• Schedule: transfer models use 5 head-only epochs (backbone frozen) then up to 30 fine-tune epochs with early stopping (patience 5); the custom CNN runs up to 20 epochs from scratch
• Loss: BCE-with-logits (binary); cross-entropy with inverse-frequency class weights (three-class)
• Checkpoint selection: lowest val loss (binary); macro-F1 (three-class)
• Seeds: 0, 1, 2

Speeds, Sizes, Times

Checkpoint sizes: custom CNN ≈1 MB, DenseNet121 ≈84 MB, ResNet18 ≈129 MB. Inference ≈13–16 ms/image on the training GPU.

Evaluation

Testing Data, Factors & Metrics

Testing Data

The dataset's held-out test/ split: 624 images (234 NORMAL, 242 BACTERIA, 148 VIRUS).

Factors

Results are reported per model, per class (NORMAL / BACTERIA / VIRUS), and as mean ± std over three seeds. Pneumonia recall is emphasised because, in screening, a missed pneumonia (false negative) is more costly than a false alarm.

Metrics

Accuracy, precision, recall (sensitivity), specificity, F1 (macro-F1 for three classes), one-vs-rest AUROC, and Expected Calibration Error (ECE). Confusion matrices and Grad-CAM overlays support qualitative analysis.

Results

Binary (NORMAL vs PNEUMONIA), test set, mean over seeds:

The binary ensemble misses only 1 of 390 pneumonia cases (recall 0.997); specificity is 0.62.

Three-class (NORMAL / BACTERIA / VIRUS), test set, patient-aware split:

The hierarchical two-stage design did not beat the flat softmax: evidence that the viral/bacterial boundary is intrinsic to the data, not a modelling artefact. Calibration: 3-class ECE ≈ 0.13 to 0.15 for single models and ≈ 0.05 for the ensemble.

Summary

Transfer learning clearly beats training from scratch. DenseNet121 is the best single model and the most interpretable; the ResNet18 + DenseNet121 ensemble is the strongest predictor on both tasks. Viral pneumonia is the hardest class throughout.

Model Examination

Grad-CAM was used on all four binary outcome types (TP/TN/FP/FN) and on the three-class outcomes. DenseNet121's heatmaps concentrate on clinically relevant lung regions (e.g. the affected lobe for bacterial consolidation), making its decisions auditable; ResNet18's are more often diffuse or centred on non-parenchymal structures even when the prediction is correct.

Environmental Impact

Carbon emissions can be estimated using the Machine Learning Impact calculator presented in Lacoste et al. (2019).

• Hardware Type: NVIDIA RTX 3060 Laptop GPU
• Hours used: ~5 hours total...

checkpoints/binary/resnet18.pt

pt · 128 MB · SHA-256 c1f4b6de43f5…0767 · Hugging Face