Chest X-ray Pneumonia Classifier with MC-Dropout Uncertainty

A DenseNet121 chest X-ray classifier (CheXNet-style backbone) that reports a per-image uncertainty score alongside its prediction, and uses that uncertainty to drive a selective-referral (risk-coverage) analysis: when the model defers its most-uncertain cases to a radiologist, AUROC on the cases it keeps goes up.

Model description

Backbone: DenseNet121, ImageNet-pretrained.
Head: Dropout(p=0.3) -> Linear(in_features, 1), kept active at inference time (Monte Carlo Dropout) so repeated stochastic forward passes give a distribution over predictions.
Task: binary classification, chest X-ray -> {NORMAL, PNEUMONIA}.
Training data: Kaggle Chest X-Ray Images (Pneumonia) dataset (Kermany et al.), 5,219 train / 19 val / 625 test images.
Training: 5 epochs, AdamW (lr=1e-4, weight_decay=1e-4), BCEWithLogitsLoss, on a single NVIDIA B200 GPU (HiPerGator).

Output

The model returns a single logit. torch.sigmoid(logit) gives the probability that the image is class index 1 (PNEUMONIA). Running mc_predict() (30 stochastic forward passes with dropout active) returns two values per image instead of one:

mean probability across passes (the prediction)
std across passes (the epistemic uncertainty score)

This model does not output segmentation masks, bounding boxes, or localization — whole-image binary classification only.

Results

Full test set: AUROC 0.967, ECE 0.320
Uncertainty tracks error: mean MC-dropout uncertainty is 0.008 on correct predictions vs. 0.030 on wrong ones (~3.7x higher).
Selective referral: AUROC on the retained set rises from 0.967 at full coverage to 0.987 when the most-uncertain 30% of cases are referred away. See referral_curve.png.

Honest limitations

ECE of 0.320 indicates the model is overconfident. This doesn't undermine the referral result (which relies on relative uncertainty ranking, not calibrated probabilities), but it matters for any real deployment claim — temperature scaling would be a natural next step.
The validation split used during training is tiny (19 images) and saturates at AUROC 1.000, so it's a weak model-selection signal; the 625-image test set is the trustworthy number.
Pneumonia detection (binary) is an easier task than multi-label nodule/mass detection (e.g. NIH ChestX-ray14); this is meant as a methodology demonstration, not a diagnostic-grade nodule detector.
MC Dropout captures epistemic uncertainty only; deep ensembles are a stronger (heavier) baseline.
No PHI is involved (public de-identified dataset), but real clinical use would require DICOM handling, de-identification, and IRB review.

How to use

import torch
from model import build_model, mc_predict  # from this repo
from data import _tf  # preprocessing transform

model = build_model(num_classes=1)
model.load_state_dict(torch.load("ckpt.pt", map_location="cpu"))

# x: a batch of preprocessed images, shape (B, 3, 224, 224)
mean_prob, uncertainty = mc_predict(model, x, n_passes=30)

Files

ckpt.pt — trained model weights (state dict).
model.py — DenseNet121 + MC-dropout head; mc_predict for sampled inference.
data.py — preprocessing transforms (grayscale->3ch, resize, ImageNet norm).
referral_curve.png — AUROC-vs-coverage plot (the headline result).

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