pulmonet-ai

PulmoNet AI

Multi-label thoracic pathology classification on the NIH ChestX-ray14 dataset, with Grad-CAM explanations and a containerised inference API.

Research and educational use only. PulmoNet AI is not a medical device, has not been clinically validated, and must not be used for diagnosis or treatment decisions.

---

What this is

An end-to-end pipeline over all 112,120 chest X-rays from 30,805 patients:

• Ingest — downloads the dataset, recovers and verifies the official identifiers, and emits a manifest with patient-level split assignments.
• Train — EfficientNet-B4 fine-tuned in PyTorch with class-weighted BCE and dynamic augmentation, sized to fit a 4 GB laptop GPU.
• Explain — Grad-CAM implemented directly against autograd (no wrapper library).
• Serve — FastAPI + Docker on AWS Lambda (torch, both endpoints), with a free static Hugging Face Space as the demo frontend.

The demo is a free static Space; the API is an AWS Lambda container serving both classification and Grad-CAM. First request after idle takes ~15 s (cold start), then ~200 ms.

pulmonet/     config, data, model, losses, gradcam, engine, inference
scripts/      ingest, train, evaluate, export_onnx
app/          FastAPI app, Lambda handler, local demo UI
space/        static frontend deployed to Hugging Face Spaces
notebooks/    kaggle_train.ipynb - same pipeline on a free cloud GPU
docker/       Dockerfile (generic) and Dockerfile.lambda (AWS)
tests/        33 tests: data contract, Grad-CAM, API, and backend equivalence
reports/      metrics.json and auroc_table.md  (generated, never hand-edited)

---

Results

See reports/metrics.json and reports/auroc_table.md. Those files are written by scripts/evaluate.py and are the only source of truth for numbers quoted about this project.

Evaluation is on the official NIH test split — 25,596 images from 2,797 patients, with no patient appearing in both training and test.

<!-- RESULTS -->

Model: nih-densenet121 at 224px.

<!-- /RESULTS -->

---

Quickstart

pip install -r requirements-train.txt
# install the torch build matching your CUDA runtime, e.g.
pip install torch --index-url https://download.pytorch.org/whl/cu118

python scripts/ingest.py --limit 200     # dry run: verifies the pipeline in ~30 s
python scripts/ingest.py                 # full ingest, ~15 min, ~1.7 GB on disk

python scripts/train.py --smoke          # 50 steps, verifies the training loop
python scripts/train.py                  # full run
python scripts/evaluate.py               # writes reports/

uvicorn app.main:app --port 7860         # http://localhost:7860
pytest

Set PULMONET_DATA to control where images and checkpoints go (default C:\pulmonet_data). Keep it outside any cloud-synced folder.

Training on a free cloud GPU

notebooks/kaggle_train.ipynb runs the same pipeline on Kaggle's free P100 (16 GB, 30 GPU-hours/week). Two reasons it is the better option than a 4 GB laptop GPU here:

• ChestX-ray14 is already hosted on Kaggle, so --source kaggle reads it in place — no download at all.
• 16 GB fits batch 32 without gradient accumulation, and nothing competes for VRAM.

Roughly 3–4× faster end to end. scripts/ingest.py --source kaggle is the only piece that differs; training, evaluation, and reporting are byte-identical.

Deployment: see DEPLOYMENT.md.

---

Design decisions

These are the choices worth understanding; each one is a place where the obvious approach is wrong.

The split is patient-level, not image-level

A patient contributes multiple X-rays. Splitting randomly by image puts near-duplicate views of the same chest on both sides of the train/test boundary, which inflates AUROC by several points and measures memorisation rather than generalisation.

This project uses the official train_val_list.txt / test_list.txt, which partition by patient — verified empirically: 28,008 training patients and 2,797 test patients, zero overlap. The validation set is then carved out of the training patients, again by patient ID (pulmonet/data.py:split_manifest), and tests/test_data.py asserts all three splits are mutually disjoint.

It is multi-label, not multi-class

One X-ray can show Effusion and Infiltration and Atelectasis simultaneously. So the model emits 14 independent sigmoids, not a softmax over 14 classes, and trains with BCEWithLogitsLoss. Probabilities do not sum to 1, and they are not supposed to.

"No Finding" is deliberately not a 15th class — it is the absence of all 14 labels, so an all-zero prediction vector already expresses it. Modelling it separately would double-count.

The class weights are capped

ChestX-ray14 prevalence ranges from Infiltration at ~17.7% to Hernia at ~0.2%. Unweighted BCE scores 99.8% "accuracy" on Hernia by always predicting negative, which is worthless — hence pos_weight = n_neg / n_pos per class.

But Hernia's raw ratio is ~500×, and feeding that in lets a single positive sample dominate the batch gradient; the loss diverges within a few hundred steps. Clipping at 10× keeps rare classes meaningfully up-weighted while bounding any one class's contribution (pulmonet/losses.py:compute_pos_weight).

The augmentation respects anatomy

Horizontal flip, ±10° rotation, mild brightness/contrast, and RandomResizedCrop at 0.85–1.0. Deliberately no vertical flip and no aggressive rotation: thoracic anatomy has a fixed orientation and the heart sits on one side, so an upside-down chest X-ray is not a sample the model will ever see at inference. Augmenting toward it spends capacity on an impossible input.

AUROC, not accuracy

With Hernia at 0.2% prevalence, accuracy is uninformative. AUROC is threshold-free and prevalence-insensitive, and it is the metric the ChestX-ray14 literature reports, so results here are directly comparable to Wang et al. (2017), which evaluate.py prints alongside.

Classes with no positive samples in a split report NaN rather than a fabricated 0.5.

Fitting EfficientNet-B4 into 4 GB of VRAM

B4's native resolution is 380px, which OOMs on an RTX 3050 Laptop at any useful batch size. The working configuration:

Batch 16 measures at 2.3 GB reserved and looks fine on an idle card — but a running browser holds ~2 GB of VRAM for compositing, leaving under 2 GB. Training then dies partway through an epoch with CUDA out of memory (or, more confusingly, CUDNN_STATUS_INTERNAL_ERROR, or CUBLAS_STATUS_NOT_INITIALIZED — cuDNN and cuBLAS both report allocation failures as opaque init errors). Batch 8 peaks near 1.25 GB and trains reliably on a machine in use.

Checkpoint selection is on validation mean AUROC, not validation loss — the pos-weighted loss is dominated by rare classes and can improve while ranking quality is flat.

Grad-CAM is written directly against autograd

pulmonet/gradcam.py is ~60 lines using a forward hook plus a tensor gradient hook on conv_head. No pytorch-grad-cam dependency, so every line is explicable.

Two non-obvious details it gets right:

• Hooks are removed in __exit__. A leaked hook accumulates on every API request.
• It forces torch.enable_grad() internally. The API wraps inference in no_grad, under which a naive implementation silently returns an all-zero map that still looks like a heatmap.

What Grad-CAM is for here: confirming the network attends to lung fields rather than to text burn-ins, scanner artifacts, or image borders — a documented shortcut-learning failure mode in chest X-ray models. It is not clinical localisation and not evidence of diagnostic validity.

Ingest verifies its own assumptions

The 300px source mirror ships (image, labels) but drops the original filenames, so patient IDs and official split membership are not directly available. Its rows are in official split-list order and its counts match exactly (86,524 / 25,596), so position i of the train split should be train_val_list.txt[i].

That is an assumption, and a wrong one would silently corrupt every metric in the project. So scripts/ingest.py re-derives each row's expected findings from Data_Entry_2017.csv and asserts they match the labels the parquet actually carries — for all 112,120 rows. Any mismatch aborts the ingest. --source nih-box fetches the original 1024px archives instead, streaming one at a time so peak disk stays ~13 GB rather than ~90 GB.

---

Known limitations

• Labels are NLP-mined from radiology reports, not expert-annotated. The NIH authors estimate >90% accuracy, so a few percent of training labels are wrong and the test set inherits the same noise. This caps achievable AUROC regardless of architecture.
• No external validation. Performance on CheXpert, MIMIC-CXR, or any clinical population is unmeasured; CXR models are known to degrade badly across scanners and sites.
• Frontal views only, and no patient metadata (age, sex, prior studies) is used.
• Not calibrated. The sigmoid outputs are ranking scores; the pos-weighting deliberately distorts them away from true posterior probabilities. Do not read them as "X% chance of disease."

---

Dataset and citation

Wang X., Peng Y., Lu L., Lu Z., Bagheri M., Summers R.M. ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases. CVPR 2017.

Dataset released by the NIH Clinical Center: https://nihcc.app.box.com/v/ChestXray-NIHCC