Metadata-Version: 2.4
Name: xrayzoo
Version: 1.0.0
Summary: A PyTorch-based model zoo for X-ray image classification across multiple medical imaging datasets.
Author-email: Subham Divakar <shubham.divakar@gmail.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/shubham10divakar/XRayZoo
Project-URL: Repository, https://github.com/shubham10divakar/XRayZoo
Project-URL: Bug Tracker, https://github.com/shubham10divakar/XRayZoo/issues
Project-URL: HuggingFace, https://huggingface.co/sdmlai
Keywords: xray,medical imaging,deep learning,pytorch,chest xray,pneumonia,covid,fracture,model zoo,pretrained models,transfer learning
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Healthcare Industry
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=1.12.0
Requires-Dist: torchvision>=0.13.0
Requires-Dist: timm>=0.6.0
Requires-Dist: Pillow>=8.0.0
Requires-Dist: numpy>=1.20.0
Provides-Extra: viz
Requires-Dist: matplotlib>=3.4.0; extra == "viz"
Provides-Extra: gradcam
Requires-Dist: grad-cam>=1.4.0; extra == "gradcam"
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: isort; extra == "dev"
Requires-Dist: flake8; extra == "dev"
Requires-Dist: twine; extra == "dev"
Requires-Dist: build; extra == "dev"
Provides-Extra: all
Requires-Dist: timm>=0.6.0; extra == "all"
Requires-Dist: matplotlib>=3.4.0; extra == "all"
Requires-Dist: grad-cam>=1.4.0; extra == "all"
Dynamic: license-file

# 🩻 XRayZoo

[![PyPI version](https://badge.fury.io/py/xrayzoo.svg)](https://pypi.org/project/xrayzoo/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
[![Python Versions](https://img.shields.io/pypi/pyversions/xrayzoo.svg)](https://pypi.org/project/xrayzoo/)
[![Framework: PyTorch](https://img.shields.io/badge/Framework-PyTorch-orange.svg)](https://pytorch.org/)

---

**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

```bash
pip install xrayzoo
```

For transformer models (ViT, Swin, etc.):
```bash
pip install xrayzoo[timm]
```

For Grad-CAM visualizations:
```bash
pip install xrayzoo[gradcam]
```

Install everything:
```bash
pip install xrayzoo[all]
```

---

## 🧬 Supported Datasets & Models

### NIH Chest X-Ray 14 (`chestxray14`)
Multilabel classification — 14 disease labels.

| Model Key | Architecture | Accuracy | AUC | Parameters |
|---|---|---|---|---|
| `densenet121_chestxray14_v1` | DenseNet121 | 82.1% | 0.841 | 8M |
| `resnet50_chestxray14_v1` | ResNet50 | 80.3% | 0.821 | 25M |
| `efficientnetb4_chestxray14_v1` | EfficientNet-B4 | 83.7% | 0.856 | 19M |
| `vit_base_chestxray14_v1` | ViT-Base-16 | 84.2% | 0.862 | 86M |

### RSNA Pneumonia Detection (`rsna_pneumonia`)
Binary classification — Normal vs Pneumonia.

| Model Key | Architecture | Accuracy | AUC |
|---|---|---|---|
| `resnet50_rsna_v1` | ResNet50 | 91.2% | 0.924 |
| `densenet121_rsna_v1` | DenseNet121 | 92.8% | 0.941 |

### COVID-19 Chest X-Ray (`covidxray`)
Multiclass — COVID-19 vs Normal vs Pneumonia.

| Model Key | Architecture | Accuracy | AUC |
|---|---|---|---|
| `resnet50_covid_v1` | ResNet50 | 95.3% | 0.978 |
| `mobilenetv2_covid_v1` | MobileNetV2 | 93.1% | 0.961 |

### Bone Fracture X-Ray (`bone_fracture`)
Multiclass — 7 fracture types.

| Model Key | Architecture | Accuracy | AUC |
|---|---|---|---|
| `densenet169_bone_v1` | DenseNet169 | 88.4% | 0.934 |

---

## ✅ Verified Working Example (Google Colab)

The following was tested and confirmed working on Google Colab:

```python
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

```python
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

```python
from xrayzoo import XRayPredictor

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

**Multiclass output:**
```json
{
    "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):**
```json
{
    "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

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

### 4. Feature Extraction

```python
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

```python
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

```python
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

```python
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

```bash
# 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`:

```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:
```bash
export XRAYZOO_CACHE_DIR=/custom/path
```
Or in Python:
```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](https://github.com/shubham10divakar/XRayZoo)  
**HuggingFace:** [huggingface.co/sdmlai](https://huggingface.co/sdmlai)

---

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