CheXNet-style DenseNet-121 for 14 chest X-ray findings

A DenseNet-121 fine-tuned on the NIH ChestX-ray14 dataset. It takes a chest X-ray image and gives a probability for each of the 14 thoracic findings below. It follows the approach of CheXNet (Rajpurkar et al., 2017).

The model was built for an academic PFA project. It is the AI part of the web app chexnet-medical-app.

[!WARNING] This model is for research and education only. It is not a medical device. Do not use it to diagnose patients or make clinical decisions. Read the limitations below first.

Model details

Classes (output order): Atelectasis, Cardiomegaly, Effusion, Infiltration, Mass, Nodule, Pneumonia, Pneumothorax, Consolidation, Edema, Emphysema, Fibrosis, Pleural_Thickening, Hernia

Files

Both weight files hold the same weights. They give identical outputs.

How to use

import torch
from PIL import Image
from torchvision import transforms
from torchvision.models import densenet121
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

CLASS_NAMES = ["Atelectasis", "Cardiomegaly", "Effusion", "Infiltration", "Mass", "Nodule",
               "Pneumonia", "Pneumothorax", "Consolidation", "Edema", "Emphysema",
               "Fibrosis", "Pleural_Thickening", "Hernia"]

model = densenet121(num_classes=14)
model.load_state_dict(load_file(hf_hub_download("TheBiggyBang/chexnet-densenet121", "model.safetensors")))
model.eval()

preprocess = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

image = Image.open("chest_xray.png").convert("RGB")
with torch.no_grad():
    probs = torch.sigmoid(model(preprocess(image).unsqueeze(0)))[0]

for name, p in sorted(zip(CLASS_NAMES, probs.tolist()), key=lambda x: -x[1]):
    print(f"{name:20s} {p:.3f}")

The training network ended in nn.Sigmoid(). That layer has no weights, so it is not stored in the files. Apply torch.sigmoid yourself, and do not use softmax: a single image can show several findings at once.

Training data

• Dataset: NIH ChestX-ray14 (Wang et al., 2017), using the Kaggle version pre-resized to 224 × 224 (nih-chest-x-ray-14-224x224-resized).
• Filtering: every image labelled "No Finding" was removed, which leaves 51,708 images that each have at least one finding.
• Labels: the Finding Labels column was converted to a 14-dim multi-hot vector.
• Split: two random splits at the image level, 80/20 then 80/20 again:

Training procedure

The published weights are from the end of epoch 30. Validation mean AUROC rose from 0.705 after epoch 1 to 0.800 after epoch 30, while training loss fell from 0.282 to 0.240. The validation AUROC logged during training is an average of per-batch AUROCs.

Evaluation

These are per-class AUROC values on the held-out test split (10,342 images), computed over the whole split:

Limitations and biases

• There is no "normal" output. "No Finding" images were removed before training, so the model never learned what a healthy X-ray looks like. Some probabilities will always come out higher than others, even for a normal image. Do not read the top class as "the patient has this disease."
• The test scores are probably too optimistic. The split was random at the image level and not grouped by patient. ChestX-ray14 has many follow-up scans of the same person, so one patient can appear in both the training and test sets. The split was also unseeded, so it cannot be reproduced exactly.
• The labels are noisy. NIH labels were extracted from radiology reports with NLP and were not read from the images by radiologists.
• Some findings are weak. Pneumonia, Infiltration and Consolidation score below 0.70 AUROC.
• The data has limited range. All images come from one U.S. hospital system and were downsampled to 224 × 224. Accuracy will likely drop on other hospitals, scanners or patient populations.
• The probabilities are not calibrated. No threshold tuning or calibration was done, and 0.5 is not a meaningful cut-off for every class.

Intended use

• In scope: education, demos, and research prototypes, such as the PFA web app this model was built for.
• Out of scope: any clinical, diagnostic or triage use, and any use where a prediction could affect a patient's care.

References

• 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.
• Rajpurkar P., Irvin J., et al. CheXNet: Radiologist-Level Pneumonia Detection on Chest X-Rays with Deep Learning. arXiv:1711.05225, 2017.
• Huang G., Liu Z., van der Maaten L., Weinberger K. Q. Densely Connected Convolutional Networks. CVPR 2017.