🩻 XRayZoo

XRayZoo is a PyTorch-based model zoo for X-ray image classification.
It provides plug-and-play pretrained models across multiple medical imaging datasets — with a unified API designed to scale across any number of datasets and architectures.

🚀 Features

• ✅ Pretrained PyTorch models — Ready-to-use architectures fine-tuned on real X-ray datasets
• 🧠 Unified API — One interface for all datasets, tasks, and architectures
• 🌐 HuggingFace-backed — Models download automatically on first use and are cached locally
• 🔬 Feature Extraction — Extract intermediate representations for research and analysis
• ⚙️ Model Introspection — Access weights, layers, and architecture details programmatically
• 🎯 Multi-task Support — Binary, multiclass, and multilabel classification in one library
• 🧩 Easily Extensible — Adding a new dataset or model requires only a JSON registry edit
• 🔧 Fine-tuning Helpers — Freeze/unfreeze backbone, save/load checkpoints
• 💻 CLI Tool — xrayzoo command for terminal-based browsing and inference
• 🔍 Advanced Search — Filter models by dataset, architecture, tag, task, or accuracy

📦 Installation

pip install xrayzoo

For transformer models (ViT, Swin, etc.):

pip install xrayzoo[timm]

For Grad-CAM visualizations:

pip install xrayzoo[gradcam]

Install everything:

pip install xrayzoo[all]

🧬 Supported Datasets & Models

NIH Chest X-Ray 14 (chestxray14)

Multilabel classification — 14 disease labels.

RSNA Pneumonia Detection (rsna_pneumonia)

Binary classification — Normal vs Pneumonia.

COVID-19 Chest X-Ray (covidxray)

Multiclass — COVID-19 vs Normal vs Pneumonia.

Bone Fracture X-Ray (bone_fracture)

Multiclass — 7 fracture types.

✅ Verified Working Example (Google Colab)

The following was tested and confirmed working on Google Colab:

from xrayzoo import XRayPredictor

predictor = XRayPredictor(model_name="swin_tiny_chestxray14_v1")
result = predictor.predict("/content/00000008_002.png")
print(result)

Output:

⬇ Downloading nih_chest_xray/swin_tiny_patch4_window7_224/swin_tiny_patch4_window7_224_NihChestXrayMclass.pth from sdmlai/xrayzoo...
  [████████████████████] 100%
✔ Saved to: /root/.xrayzoo/models/chestxray14/nih_chest_xray/...

=======================================================
  Model    : swin_tiny_chestxray14_v1
  Dataset  : chestxray14
  Task     : multilabel
  Classes  : 14
  Device   : cpu
  Accuracy : 99.7%
=======================================================

{
  'model': 'swin_tiny_chestxray14_v1',
  'dataset': 'chestxray14',
  'task': 'multilabel',
  'threshold': 0.5,
  'predictions': [
    {'label': 'Cardiomegaly',      'confidence': 0.082755, 'positive': False},
    {'label': 'Atelectasis',       'confidence': 0.061377, 'positive': False},
    {'label': 'Fibrosis',          'confidence': 0.059574, 'positive': False},
    {'label': 'Pleural_Thickening','confidence': 0.055506, 'positive': False},
    {'label': 'Infiltration',      'confidence': 0.049657, 'positive': False},
    {'label': 'Nodule',            'confidence': 0.041546, 'positive': False},
    {'label': 'Effusion',          'confidence': 0.027421, 'positive': False},
    {'label': 'Emphysema',         'confidence': 0.019398, 'positive': False},
    {'label': 'Hernia',            'confidence': 0.01718,  'positive': False},
    {'label': 'Mass',              'confidence': 0.009219, 'positive': False},
    {'label': 'Pneumothorax',      'confidence': 0.008874, 'positive': False},
    {'label': 'Consolidation',     'confidence': 0.007585, 'positive': False},
    {'label': 'Pneumonia',         'confidence': 0.005516, 'positive': False},
    {'label': 'Edema',             'confidence': 0.000554, 'positive': False}
  ],
  'positive_labels': []
}

positive_labels: [] means no disease was predicted above the 0.5 threshold — the model classified this image as No Finding (normal chest X-ray). Lower the threshold to see confidence scores for each condition.

Note: Model weights are downloaded automatically on first use (~110 MB) and cached at ~/.xrayzoo/models/.

🔬 How to Use

1. Explore the Zoo

from xrayzoo import zoo

# Summary of everything
zoo.summary()

# List all datasets
zoo.list_datasets()

# List models for a specific dataset
zoo.list_models(dataset="chestxray14")

# Detailed model info
zoo.info("densenet121_chestxray14_v1")

# Leaderboard (sorted by accuracy or AUC)
zoo.leaderboard(dataset="chestxray14", metric="auc")

# Compare models side by side
zoo.compare([
    "densenet121_chestxray14_v1",
    "vit_base_chestxray14_v1",
    "efficientnetb4_chestxray14_v1"
])

# Advanced search
zoo.search(tag="transformer", min_accuracy=83.0)
zoo.search(task="binary", architecture="densenet121")

2. Predict

from xrayzoo import XRayPredictor

predictor = XRayPredictor(model_name="densenet121_chestxray14_v1")
result = predictor.predict("path/to/xray.jpg")
print(result)

Multiclass output:

{
    "model": "resnet50_covid_v1",
    "dataset": "covidxray",
    "task": "multiclass",
    "label": "COVID-19",
    "confidence": 0.962,
    "top_k": [
        {"label": "COVID-19", "confidence": 0.962},
        {"label": "Pneumonia", "confidence": 0.031},
        {"label": "Normal", "confidence": 0.007}
    ]
}

Multilabel output (ChestX-ray14):

{
    "model": "densenet121_chestxray14_v1",
    "dataset": "chestxray14",
    "task": "multilabel",
    "threshold": 0.5,
    "positive_labels": ["Atelectasis", "Effusion"],
    "predictions": [
        {"label": "Atelectasis", "confidence": 0.842, "positive": true},
        {"label": "Effusion",    "confidence": 0.713, "positive": true},
        {"label": "Pneumonia",   "confidence": 0.124, "positive": false}
    ]
}

3. Batch Prediction

predictor = XRayPredictor(model_name="densenet121_rsna_v1")
results = predictor.predict_batch([
    "xray1.jpg",
    "xray2.jpg",
    "xray3.jpg"
])

4. Feature Extraction

predictor = XRayPredictor(model_name="densenet121_chestxray14_v1")

# Auto-extract from penultimate layer
features = predictor.extract_features("xray.jpg")
# → numpy array of shape (1024,)

# Extract from a specific layer
features = predictor.extract_features("xray.jpg", layer_name="features.denseblock3")

5. Model Introspection

predictor = XRayPredictor(model_name="densenet121_chestxray14_v1")

# Get full PyTorch model (for fine-tuning, ONNX export, etc.)
model = predictor.get_model()

# Get weights
weights = predictor.get_weights()              # state_dict
weights_info = predictor.get_weights_info()    # shape info

# List all layers
layers = predictor.list_layers()
for layer in layers[:5]:
    print(layer)

# Count parameters
params = predictor.count_parameters()
print(f"Total: {params['total']:,}  Trainable: {params['trainable']:,}")

6. Fine-tuning

import torch
from xrayzoo import XRayPredictor

predictor = XRayPredictor(model_name="densenet121_chestxray14_v1")

# Freeze backbone, only train classifier head
predictor.freeze_backbone()

# Or unfreeze everything
predictor.unfreeze_all()

# Get the model for your training loop
model = predictor.get_model()
optimizer = torch.optim.Adam(
    filter(lambda p: p.requires_grad, model.parameters()),
    lr=1e-4
)

# Save your fine-tuned checkpoint
predictor.save_checkpoint("my_finetuned.pth", extra_meta={"notes": "my experiment"})

7. Device Control

import torch

# Auto-detect (CUDA → MPS → CPU)
predictor = XRayPredictor("densenet121_chestxray14_v1")

# Force specific device
predictor = XRayPredictor(
    "densenet121_chestxray14_v1",
    device=torch.device("cuda:0")
)

🖥️ CLI Usage

# Summary
xrayzoo summary

# Browse
xrayzoo datasets
xrayzoo models
xrayzoo models --dataset chestxray14

# Model info
xrayzoo info densenet121_chestxray14_v1

# Leaderboard
xrayzoo leaderboard --dataset chestxray14 --metric auc

# Search
xrayzoo search --tag transformer --min-accuracy 83.0

# Predict
xrayzoo predict xray.jpg --model densenet121_chestxray14_v1 --top-k 3

# Cache management
xrayzoo cache list
xrayzoo cache clear
xrayzoo cache clear --dataset chestxray14

🧩 How to Add a New Dataset (Scalable Design)

XRayZoo is designed so that adding a new dataset requires zero Python code changes — only a JSON edit.

1. Train your model on the new dataset and upload the .pth file + labels JSON to HuggingFace.

2. Add a new entry to xrayzoo/model_registry.json:

{
  "datasets": {
    "your_new_dataset": {
      "name": "Your Dataset Name",
      "description": "Dataset description",
      "num_classes": 5,
      "task": "multiclass",
      "input_size": [224, 224],
      "labels": ["Class A", "Class B", "Class C", "Class D", "Class E"],
      "models": {
        "resnet50_your_dataset_v1": {
          "description": "ResNet50 trained on your dataset",
          "architecture": "resnet50",
          "input_size": [224, 224],
          "accuracy": 92.0,
          "auc": 0.96,
          "parameters": "25M",
          "hf_repo": "sdmlai/xrayzoo",
          "model_filename": "resnet50_your_dataset_v1.pth",
          "labels_filename": "your_dataset_labels.json",
          "preprocessing": "chestxray_standard",
          "task": "multiclass",
          "tags": ["cnn", "resnet", "your-dataset", "multiclass"]
        }
      }
    }
  }
}

3. That's it. The model is now available across the full API: zoo.list_models(), XRayPredictor(...), leaderboard, search, CLI, etc.

📁 Project Structure

xrayzoo/
├── xrayzoo/
│   ├── __init__.py          # Public API exports
│   ├── model_registry.json  # All datasets & models (the scalable data layer)
│   ├── registry.py          # ModelRegistry class — reads registry JSON
│   ├── predictor.py         # XRayPredictor — main inference class
│   ├── zoo.py               # Zoo — discovery & display interface
│   ├── downloader.py        # HuggingFace download + local cache
│   ├── model_loader.py      # Architecture builder + weight loader
│   ├── preprocessing.py     # Image transform pipelines
│   ├── cli.py               # `xrayzoo` CLI tool
│   └── exceptions.py        # Custom exceptions
├── tests/
│   └── test_xrayzoo.py
├── setup.py
├── pyproject.toml
├── MANIFEST.in
└── README.md

🗃️ Local Cache

Models are cached in ~/.xrayzoo/models/<dataset>/ after first download.

Override the cache directory:

export XRAYZOO_CACHE_DIR=/custom/path

Or in Python:

predictor = XRayPredictor("densenet121_chestxray14_v1", cache_dir="/custom/path")

📝 License

MIT License — free for academic and commercial use.

🤝 Contributing

Contributions are welcome! You can help by:

• 🧠 Training and uploading new models to HuggingFace
• 📊 Adding new X-ray datasets to the registry
• ⚙️ Adding new architectures to model_loader.py
• 🧪 Writing tests and benchmarks
• 📚 Improving documentation

📫 Contact

Author: Subham Divakar
Email: shubham.divakar@gmail.com
GitHub: shubham10divakar/XRayZoo
HuggingFace: huggingface.co/sdmlai

"Bringing pretrained X-ray intelligence to every researcher — one model at a time." 🩻