Pneumonia Classifier (Custom CNN)

A lightweight custom CNN for binary classification of chest X-rays (Normal vs Pneumonia) with Grad-CAM explainability.

Model Details

Architecture: Custom CNN (Net) — 9 Conv2d layers + Global Average Pooling
Input: RGB chest X-ray images, resized to 224×224
Output: Binary classification (Normal / Pneumonia) with softmax probabilities
Parameters: ~35K (lightweight, edge-friendly)
Framework: PyTorch

Demo

Try the live demo: Pneumonia AI Diagnostic Suite

Performance

Metric	Score
Accuracy	98%
Pneumonia Precision	97%
Pneumonia Recall	100%
Normal Precision	100%
Normal Recall	96%
Macro F1-Score	0.98
ROC AUC	1.00

Evaluated on 57 test samples with optimized threshold (0.1) for high sensitivity.

Cross-Validation (5-Fold Stratified)

Fold	Accuracy	F1-Score
1	92.86%	0.927
2	97.62%	0.976
3	92.86%	0.933
4	95.24%	0.950
5	97.62%	0.977
Mean	95.24%	0.953

Training

Optimizer: SGD with Momentum (0.8)
Loss: Negative Log Likelihood (NLL) on log_softmax output
Epochs: 18 (early stopping, patience=5)
Best Training Accuracy: 98.25%
Data Augmentation: Baseline (no augmentation) — the unaugmented model performed better (95.24% CV vs 92.38% with heavy augmentation)

Quantization (INT8)

	FP32	INT8	Improvement
Size	243 KB	52 KB	4.64x smaller
Latency	108 ms	27 ms	3.99x faster

Usage

import torch
from huggingface_hub import hf_hub_download
from pneumonia_classifier.ml.model.arch import Net

# Download model
model_path = hf_hub_download(
    repo_id="24f2004275/pneumonia_classifier",
    filename="pneumonia_classifier_cnn_uza7heywpgthvahb.pt"
)

# Load model
model = Net()
model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=False))
model.eval()

# Inference
from torchvision import transforms

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

# image = Image.open("chest_xray.jpg").convert("RGB")
# tensor = transform(image).unsqueeze(0)
# with torch.no_grad():
#     output = model(tensor)
#     probs = torch.exp(output)
#     prediction = "Pneumonia" if probs.argmax() == 1 else "Normal"

Available Models

File	Description
`pneumonia_classifier_cnn_uza7heywpgthvahb.pt`	Primary model (FP32, best accuracy)
`pneumonia_classifier_cnn_int8_tgzwzsqwqw54dahb.pt`	INT8 quantized variant
`pneumonia_classifier_cnn_int8_ra32tviwyo4u3ahb.pt`	INT8 quantized variant (alt)
`pneumonia_classifier_aug_baseline_no_augmentation_zvyot3qwqgnsfahb.pt`	Baseline (no augmentation)
`pneumonia_classifier_aug_augmented_heavy_wypxe3qwqsx75ahb.pt`	Heavy augmentation variant

Explainability

This model integrates Grad-CAM (Gradient-weighted Class Activation Mapping) to visualize which regions of the chest X-ray the model focuses on for its predictions. The heatmap highlights areas of radiographical density associated with pneumonia.

License

MIT License

Citation

@misc{pneumonia_classifier,
  title={Pneumonia Detection from Chest X-Rays using Custom CNN with Grad-CAM},
  author={Ayush Dubey},
  year={2024},
  url={https://huggingface.co/24f2004275/pneumonia_classifier}
}

Downloads last month: -; Downloads are not tracked for this model. How to track