sandy45/ChestViT-Explainable-XRay-AI
0
Multi-Task Vision Transformer for Chest X-Ray Disease Classification & Explainability
<div align="center">
ViT-Base-16 · 14-Disease Multi-Label Classification · Attention Rollout XAI
Fine-tuned on NIH ChestX-ray14 (112,120 frontal chest X-rays)
</div>
🎯 Overview
This project implements an explainable medical AI system for automated chest X-ray analysis. Unlike standard classification models, this system simultaneously:
- Predicts 14 diseases in parallel (multi-label, not multi-class)
- Shows WHERE in the X-ray the model is looking via Attention Rollout
- Compares against published NIH baselines (AUC-ROC per class)
- Serves a live demo via a Gradio dashboard
The core insight: Vision Transformers (ViTs) divide images into 16×16 patches and route information through 12 attention layers. Attention Rollout traces this information flow back to the input patches — telling us exactly which lung regions drove each disease prediction.
🏗 Architecture
Chest X-Ray Input (PNG, 1024×1024)
│
▼
┌───────────────────────┐
│ CLAHE Preprocessing │ ← Contrast Limited Adaptive Histogram Equalization
│ + Albumentations │ ← Radiologically-realistic augmentation
└───────────────────────┘
│ 224×224×3
▼
┌───────────────────────────────────────────┐
│ ViT-Base-16 │
│ (google/vit-base-patch16-224-in21k) │
│ │
│ ┌─────────────────────────────────────┐ │
│ │ 196 Patches (14×14 grid, 16px/ea) │ │
│ │ + [CLS] token = 197 total tokens │ │
│ └─────────────────────────────────────┘ │
│ ↓ │
│ 12 Transformer Layers │
│ (12 heads × 64 dim = 768 hidden dim) │
│ ↓ │
│ [CLS] token → Dropout → Linear(768→14) │
└───────────────────────────────────────────┘
│
├─── Logits → Sigmoid → 14 disease probabilities
│
└─── Attention weights (12 layers × 12 heads)
│
▼
Attention Rollout Algorithm
(14×14 patch attention map)
│
▼
224×224 heatmap overlay on X-ray📊 Results
Results will populate after training. NIH baseline from Wang et al. (2017).
🚀 Quick Start
1. Install Dependencies
# Create virtual environment
python -m venv venv
venv\Scripts\activate # Windows
# source venv/bin/activate # Linux/macOS
# Install dependencies
pip install -r requirements.txt2. Download Dataset
# First: set up Kaggle API credentials
# 1. Go to https://www.kaggle.com → Account → Settings → Create New API Token
# 2. Place kaggle.json at: C:\Users\<YourName>\.kaggle\kaggle.json
# Then download NIH ChestX-ray14 (~42 GB)
python data/download.py3. Run Unit Tests
python -m pytest tests/ -v4. Train the Model
# Full training (5 epochs, ~8-12 hours on RTX 3050)
python training/train.py
# Quick smoke test (20% of data)
# Edit config/config.yaml → dataset.train_fraction: 0.2
python training/train.pyMonitor training in real-time:
mlflow ui --backend-store-uri ./experiments/mlflow
# Open http://localhost:50005. Launch Demo
# With trained model
python app/gradio_app.py
# DEMO MODE (random weights, for UI preview only)
set DEMO_MODE=1 # Windows
python app/gradio_app.py📁 Project Structure
.
├── config/
│ └── config.yaml # All hyperparameters and paths
├── data/
│ ├── download.py # Kaggle API dataset download
│ ├── preprocessing.py # CLAHE + Albumentations pipeline
│ ├── dataset.py # ChestXrayDataset + DataLoaders
│ └── raw/ # Downloaded dataset (not in git)
├── models/
│ └── vit_model.py # ViT-Base-16 with multi-label head
├── explainability/
│ └── attention_rollout.py # Attention Rollout algorithm
├── training/
│ ├── losses.py # Weighted BCE + Focal Loss
│ ├── train.py # Training loop (mixed precision, MLflow)
│ └── evaluate.py # AUC-ROC per class, ROC plots
├── app/
│ └── gradio_app.py # Gradio dashboard
├── tests/
│ └── test_modules.py # Unit tests (no dataset required)
├── checkpoints/ # Saved model weights (not in git)
├── results/ # ROC curves, metrics CSV
├── experiments/
│ └── mlflow/ # MLflow tracking database
├── config_loader.py # YAML config loader
└── requirements.txt⚙️ Configuration
All settings are in `config/config.yaml`. Key RTX 3050 settings:
training:
batch_size: 8 # Fits in 4 GB VRAM
gradient_accumulation_steps: 4 # Effective batch = 32
mixed_precision: true # fp16 — mandatory for 4 GB VRAM
num_epochs: 5
model:
name: "google/vit-base-patch16-224-in21k"
gradient_checkpointing: true # Saves ~30% VRAM
dataset:
train_fraction: 1.0 # Set 0.2 for quick smoke test🔥 Attention Rollout: Why Not Grad-CAM?
Attention Rollout (Abnar & Zuidema, 2020) is mathematically derived from the transformer's own attention mechanism, making it the correct tool for ViT explainability.
🏥 Clinical Context
⚠️ This is a research/educational project, NOT a medical device. Results should not be used for clinical diagnosis without radiologist review.
The NIH ChestX-ray14 dataset has known limitations (Rajpurkar et al., 2018, and others). AUC-ROC is the clinically relevant metric because:
- Accuracy is misleading with class imbalance (>53% "No Finding")
- AUC measures discriminative ability across all thresholds
- Radiologists can set their own confidence threshold per clinical context
📚 References
- Wang, X. et al. (2017). ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks. CVPR.
- Dosovitskiy, A. et al. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR.
- Abnar, S. & Zuidema, W. (2020). Quantifying Attention Flow in Transformers. arXiv:2005.00928.
- Rajpurkar, P. et al. (2017). CheXNet: Radiologist-Level Pneumonia Detection on Chest X-Rays with Deep Learning. arXiv:1711.05225.