chest-xray-multilabel-classifier

Multi-Label Chest X-Ray Classifier

ResNet-50 | PyTorch | NIH ChestX-ray14 | Grad-CAM

An end-to-end medical computer vision project using an ImageNet-pretrained ResNet-50 to classify 14 thoracic conditions from chest X-ray images. The model was trained on the NIH ChestX-ray14 dataset with patient-level data splitting, class-weighted binary cross-entropy, and experiment tracking through Weights & Biases.

Held-Out Test Mean ROC-AUC: 0.823

Live Interactive Demo | GitHub Repository

Model Overview

Supported Conditions

The model generates independent prediction scores for 14 thoracic conditions:

1. Atelectasis
2. Cardiomegaly
3. Effusion
4. Infiltration
5. Mass
6. Nodule
7. Pneumonia
8. Pneumothorax
9. Consolidation
10. Edema
11. Emphysema
12. Fibrosis
13. Pleural Thickening
14. Hernia

Because this is a multi-label classification task, multiple conditions may be associated with the same X-ray.

The model produces 14 logits, converted into scores using the sigmoid function. These scores are not clinically calibrated probabilities.

Model Architecture

The classifier uses an ImageNet-pretrained ResNet-50 backbone with a custom classification head.

Chest X-Ray (224 × 224 × 3)
             |
     ResNet-50 Backbone
             |
     2048-D Feature Vector
             |
       Linear(2048, 512)
             |
            ReLU
             |
        Dropout(0.3)
             |
        Linear(512, 14)
             |
       14 Output Logits
             |
     Sigmoid (Inference)
             |
     14 Prediction Scores

The model is implemented in PyTorch using the custom ChestXRayNet class.

Dataset and Preprocessing

The model was trained on the NIH ChestX-ray14 dataset, containing 112,120 frontal chest X-ray images from 30,805 patients.

Preprocessing includes:

• RGB image conversion
• Resizing to 224 × 224 pixels
• ImageNet mean and standard deviation normalization
• Training-time horizontal flipping, brightness/contrast adjustments, and geometric augmentation

Patient-Level Data Splitting

To prevent patient-level data leakage, the dataset was split by patient identifier rather than by individual image.

The split uses a 70/15/15 patient-level allocation with random seed 42 and zero patient overlap across partitions.

Images labeled No Finding were retained as negative examples for all 14 conditions.

Training Methodology

Training was performed on Kaggle using an NVIDIA Tesla T4 GPU.

Class Imbalance Handling

The NIH ChestX-ray14 dataset contains substantial label imbalance, particularly for rare conditions such as Hernia and Pneumonia.

To address this, positive class weights were calculated from the training set:

positive_weight = negative_samples / positive_samples

These weights were incorporated into BCEWithLogitsLoss to increase the contribution of positive examples from underrepresented classes.

Training and validation metrics were tracked using Weights & Biases.

Evaluation Results

The best checkpoint was selected at epoch 7 using validation mean ROC-AUC.

• Best Validation Mean ROC-AUC: 0.8315
• Held-Out Test Mean ROC-AUC: 0.8225
• Test Images: 16,491
• Test Patients: 4,621

Per-Class Test ROC-AUC

ROC-AUC measures the model's ability to distinguish positive and negative examples across classification thresholds. It does not establish clinical diagnostic accuracy, calibration, sensitivity, or specificity at a particular operating threshold.

Model Inference

The trained model checkpoint is available in this repository as best_model.pth.

Inference requires the matching ChestXRayNet architecture and preprocessing pipeline.

For an interactive inference experience, visit the Hugging Face Space.

The deployed application allows users to:

• Upload a chest X-ray image
• Generate prediction scores for all 14 conditions
• View conditions ranked by model score
• Generate a Grad-CAM heatmap for the highest-scoring condition
• Visualize the heatmap overlaid on the original image

Model Interpretability

Grad-CAM (Gradient-weighted Class Activation Mapping) is integrated into the model to visualize image regions contributing to a selected class output.

The implementation uses gradients and activations from the final ResNet-50 convolutional block to generate spatial heatmaps.

Grad-CAM visualizations are interpretability aids, not validated medical localization results. Highlighted regions do not necessarily correspond to clinically confirmed pathology.

Intended Use

This model is intended for:

• Educational and research exploration of medical computer vision
• Multi-label classification experimentation
• Transfer learning and class imbalance studies
• Model interpretability demonstrations
• Machine learning engineering portfolio evaluation

Limitations and Ethical Considerations

• The model was evaluated on a held-out partition of NIH ChestX-ray14, not an independent external clinical dataset.
• NIH ChestX-ray14 labels are derived from radiology reports and may contain annotation noise.
• Performance may vary across hospitals, imaging equipment, patient populations, and acquisition conditions.
• Rare conditions have relatively few positive examples, limiting the stability of their performance estimates.
• Model scores have not been clinically calibrated.
• No prospective clinical evaluation or regulatory validation has been performed.
• Grad-CAM heatmaps should not be interpreted as definitive pathology localization.
• The model should not be used for diagnosis, treatment decisions, or autonomous medical screening.

This model is a research and educational prototype, not a clinically validated medical device.

Project Resources

• Interactive Hugging Face Demo
• Source Code and Training Pipeline

Author

Samer Ahmed

Developed as a medical computer vision and deep learning engineering project.