Skip to content

Repository files navigation

xvr: X-ray to Volume Registration

docs tests Paper shield License: MIT Hugging Face Hugging Face uv

A PyTorch package for training patient-specific 2D/3D registration models in 5 minutes.

image

Highlights

  • 🚀 A single CLI/API for training models and registering clinical data
  • ⚡️ 100x faster patient-specific model training than DiffPose
  • 📐 Submillimeter registration accuracy with new image-similarity metrics
  • 🩺 Human-interpretable pose parameters for training your own models
  • 🐍 Pure Python/PyTorch implementation
  • 🖥️ Supports macOS, Linux, and Windows

xvr is built upon DiffDRR, the differentiable X-ray renderer.

Installation

Install the Python API and CLI (should take ~5 min if installing PyTorch with CUDA):

pip install git+https://github.com/eigenvivek/xvr.git

Verify the installation version (should match the latest release on GitHub):

xvr --version

CLI Usage

xvr provides a command-line interface for training/finetuning pose regression models and registering clinical data with gradient-based iterative optimization with trained models. It is designed to be modular and extensible, allowing users to easily train models on new datasets and anatomical structures without any manual annotations. Full documentation is available here.

$ xvr --help

Usage: xvr COMMAND

Commands:
register     Use gradient-based optimization to register XRAY to a CT/MR.
restart      Restart model training from a checkpoint.
train        Train a pose regression model.
--help -h    Display this message and exit.
--version    Display application version.

Development

xvr is built using uv, an extremely fast Python project manager.

If you want to modify xvr (e.g., adding different loss functions, network architectures, etc.), uv makes it easy to set up a development environment:

# Download xvr
git clone https://github.com/eigenvivek/xvr && cd xvr

# Install uv and build the environment with all dev requirements
curl -LsSf https://astral.sh/uv/install.sh | sh
uv sync --all-groups

# Install pre-commit hooks locally
uv tool install prek
uvx prek install -f

To verify your virtual environment, you can run

uv run xvr --version

Alternatively, you can directly use the virtual environment that uv creates:

source .venv/bin/activate
xvr --version

xvr's pre-commit hooks automatically take care of things like linting and formatting, so hack away! All PRs are welcome.

Experiments

Reproducing the paper's registration results requires a CUDA GPU. First, build the environment with uv:

git clone https://github.com/eigenvivek/xvr.git && cd xvr
uv sync --all-groups

Then download the pretrained models (3.7 GB) and datasets (4.8 GB) from HuggingFace:

uvx hf download eigenvivek/xvr      --repo-type model   --local-dir experiments/models/
uvx hf download eigenvivek/xvr-data --repo-type dataset --local-dir experiments/data/

Registration runs three datasets (DeepFluoro, Femur, Ljubljana) × three initializations (de novo, finetuned, foundation) as nine SLURM array jobs:

./experiments/run.sh register

The scripts are in experiments/scripts/{dataset}/register/. Four #SBATCH directives are cluster-specific: update --partition, --qos, --account, and --gres to match your platform. Metrics were computed on an NVIDIA RTX 6000 Ada with PyTorch 2.10.

If you don't have SLURM, you can run the subjects in series by manually supplying the array index:

for m in de_novo finetuned foundation; do
    for i in $(seq 1 6);  do SLURM_ARRAY_TASK_ID=$i bash experiments/scripts/deepfluoro/register/$m.sh; done
    for i in $(seq 1 5);  do SLURM_ARRAY_TASK_ID=$i bash experiments/scripts/femur/register/$m.sh;      done
    for i in $(seq 1 10); do SLURM_ARRAY_TASK_ID=$i bash experiments/scripts/ljubljana/register/$m.sh;  done
done

Once every job has finished, score the results:

./experiments/run.sh evaluate

This writes experiments/results/registration.csv, rebuilt from scratch on each run, with one row per x-ray per pose (init and final) recording mPE, mRPE, mTRE, dGeo, the final NCC, and runtime.

About

Training patient-specific 2D/3D registration models in 5 minutes

Topics

Resources

Stars

73 stars

Watchers

1 watching

Forks

Releases

Packages

Used by

Contributors

Languages