File size: 2,143 Bytes
b119597 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | ---
language: en
license: cc-by-4.0
tags:
- image-classification
- medical-imaging
- chest-xray
- pneumonia
- efficientnet
- gradcam
---
# 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](https://huggingface.co/datasets/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.
|