Migrate from diffdrr to nanodrr

.github/workflows/tests.yml

@@ -0,0 +1,22 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v5
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Run tests
+        run: uv run pytest

.pre-commit-config.yaml

@@ -1,23 +1,24 @@
 repos:
   - repo: https://github.com/abravalheri/validate-pyproject
-    rev: v0.24.1
+    rev: v0.25
     hooks:
       - id: validate-pyproject
 
   - repo: https://github.com/crate-ci/typos
-    rev: v1.30.0
+    rev: v1.43.4
     hooks:
       - id: typos
         args: ["--config", ".typos.toml"]
 
   - repo: https://github.com/astral-sh/uv-pre-commit
-    rev: 0.8.17
+    rev: 0.10.1
     hooks:
       - id: uv-lock
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.3
+    rev: v0.15.0
     hooks:
       - id: ruff-check
         args: [ --fix ]
       - id: ruff-format
+        exclude: ^src/xvr/cli/configs/(register|train)\.py$

README.md

@@ -101,6 +101,7 @@ Data options:
   --patch_size TEXT               Optional random crop size (e.g., 'h,w,d'); if None, return entire volume
   --num_workers INTEGER           Number of subprocesses to use in the dataloader  [default: 4]
   --pin_memory                    Copy volumes from the dataloader into CUDA pinned memory before returning
+  --sample_weights PATH           Probability for sampling each volume in `volpath`
 
 Sampling options:
   --r1 <FLOAT FLOAT>...           Range for primary angle (in degrees)  [required]
@@ -115,7 +116,6 @@ Renderer options:
   --sdd FLOAT                     Source-to-detector distance (in millimeters)  [required]
   --height INTEGER                DRR height (in pixels)  [required]
   --delx FLOAT                    DRR pixel size (in millimeters / pixel)  [required]
-  --renderer [siddon|trilinear]   Rendering equation  [default: trilinear]
   --orientation [AP|PA]           Orientation of CT volumes  [default: AP]
   --reverse_x_axis                Enable to obey radiologic convention (e.g., heart on right)
 
@@ -127,13 +127,14 @@ Model options:
   --convention TEXT               If `parameterization='euler_angles'`, specify order  [default: ZXY]
   --unit_conversion_factor FLOAT  Scale factor for translation prediction (e.g., from m to mm)  [default: 1000.0]
   --p_augmentation FLOAT          Base probability of image augmentations during training  [default: 0.333]
+  --use_compile                   Compile forward pass with `max-autotune-no-cudagraphs`
+  --use_bf16                      Run all ops in bf16
 
 Optimizer options:
   --lr FLOAT                      Maximum learning rate  [default: 0.0002]
   --weight_ncc FLOAT              Weight on mNCC loss term  [default: 1.0]
   --weight_geo FLOAT              Weight on geodesic loss term  [default: 0.01]
   --weight_dice FLOAT             Weight on Dice loss term  [default: 1.0]
-  --weight_mvc FLOAT              Weight on multiview consistency loss term  [default: 0]
   --n_total_itrs INTEGER          Number of iterations for training the model  [default: 1000000]
   --n_warmup_itrs INTEGER         Number of iterations for warming up the learning rate  [default: 1000]
   --n_grad_accum_itrs INTEGER     Number of iterations for gradient accumulation  [default: 4]

docs/index.md

@@ -1,5 +1,12 @@
 # X-ray to Volume Registration
 
+[![Paper shield](https://img.shields.io/badge/arXiv-2503.16309-red.svg)](https://arxiv.org/abs/2503.16309)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://github.com/eigenvivek/xvr/blob/main/LICENSE)
+<a href="https://colab.research.google.com/drive/1K9lBPxcLh55mr8o50Y7aHkjzjEWKPCrM?usp=sharing"><img alt="Colab" src="https://colab.research.google.com/assets/colab-badge.svg"></a>
+<a href="https://huggingface.co/eigenvivek/xvr/tree/main" target="_blank"><img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-ffc107?color=ffc107&logoColor=white"/></a>
+<a href="https://huggingface.co/datasets/eigenvivek/xvr-data/tree/main" target="_blank"><img alt="Hugging Face" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Data-ffc107?color=ffc107&logoColor=white"/></a>
+[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)
+
 ***Training patient-specific 2D/3D registration models in 5 minutes***
 
 - 🚀 A single CLI/API for training models and registering clinical data

pyproject.toml

@@ -10,35 +10,55 @@ authors = [
 requires-python = ">=3.10"
 dependencies = [
     "antspyx>=0.5.4",
-    "click>=8.1.8",
-    "diffdrr>=0.6.0",
-    "pydicom[pixeldata]>=3.0.1",
-    "wandb>=0.21.1",
+    "cyclopts>=4.6.0",
+    "jaxtyping>=0.3.0",
+    "nanodrr[all]>=0.1.4",
+    "pydicom[pixeldata]>=3.0.0",
+    "timm>=1.0.0",
+    "wandb>=0.22.0",
 ]
 
 [project.scripts]
-xvr = "xvr.cli:cli"
+xvr = "xvr.cli:main"
 
 [project.urls]
 Repository = "https://github.com/eigenvivek/xvr"
 Issues = "https://github.com/eigenvivek/xvr/issues"
 
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore:builtin type Swig.*has no __module__:DeprecationWarning",
+]
+
 [tool.ruff]
+line-length = 100
 extend-exclude = ["*.ipynb"]
+lint.extend-select = ["I"]
+lint.ignore = [
+    "F722",  # Forward annotation false positive from jaxtyping
+    "F821",  # Forward annotation false positive from jaxtyping
+]
+src = ["src"]
 
-[tool.ruff.lint]
-extend-select = ["I"]
+[tool.ruff.lint.per-file-ignores]
+"src/xvr/cli/configs/register.py" = ["E501"]
+"src/xvr/cli/configs/train.py" = ["E501"]
 
 [tool.ruff.lint.isort]
+known-first-party = ["xvr"]
 known-third-party = ["wandb"]
 
+[tool.uv.sources]
+nanodrr = { git = "https://github.com/eigenvivek/nanodrr" }
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [dependency-groups]
 dev = [
     "prek>=0.3.2",
+    "pytest>=9.0.2",
 ]
 docs = [
     "black>=26.1.0",

src/xvr/__main__.py

@@ -0,0 +1,3 @@
+from xvr.cli import main
+
+main()

src/xvr/cli/__init__.py

@@ -1,3 +1,16 @@
-from .cli import cli
+from cyclopts import App
 
-__all__ = ["cli"]
+from .register import register
+from .restart import restart
+from .train import train
+
+xvr = App()
+
+
+xvr.command(train)
+xvr.command(restart)
+xvr.command(register)
+
+
+def main():
+    xvr()