AurevinP's picture
Upload README.md with huggingface_hub
b119597 verified
metadata
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

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.