AurevinP/pneumonia-classifier-effnetb0

Pneumonia Classification Model (EfficientNet-B0)

Binary classifier for chest X-ray images: Pneumonia vs Normal.

Architecture

• Model: EfficientNet-B0 (timm) — 4.0M parameters
• Input: 224×224 RGB (grayscale CXR → 3-channel)
• Pretraining: ImageNet
• Head: 2-class linear classifier

Dataset

• Source: hf-vision/chest-xray-pneumonia
• Splits: Train (5,216) | Validation (16) | Test (624)
• Class imbalance: ~1:2.9 (Normal:Pneumonia in train)

Training Recipe

• Reproducibility: Fixed seed = 42, deterministic mode
• Augmentation: RandomHorizontalFlip, RandomRotation(15°), ColorJitter(brightness/contrast 0.05)
• Normalization: ImageNet stats
• Loss: Weighted CrossEntropy (inverse class frequency)
• Sampler: WeightedRandomSampler to balance batches
• Optimizer: AdamW (lr=1e-4, weight_decay=1e-4)
• Epochs: 5 (stratified 200+200 subset for balanced training)

Test Performance

Metric	Score
Accuracy	0.8125
Precision (Pneumonia)	0.7910
Recall (Pneumonia)	0.9513
F1-Score	0.8638
ROC-AUC	0.9037

Explainability

Grad-CAM heatmaps are included in gradcam/ to visualize regions influencing predictions.

Files

• model.pt — Trained model checkpoint (state_dict + config + results)
• results.json — Detailed metrics and class distribution
• cm.png — Confusion matrix
• roc.png — ROC curve
• gradcam/*.png — Grad-CAM overlays

Limitations

• Trained on a small stratified subset (400 images) due to compute constraints; full-dataset training would further improve generalization.
• Validation set is very small (16 images); results may have high variance.
• No clinical validation performed; intended for research/educational use only.
• Class imbalance in the original dataset was addressed via sampling but may not fully represent real-world prevalence.