nanodrr 0.1.0__tar.gz

nanodrr-0.1.0/PKG-INFO

@@ -0,0 +1,56 @@
+Metadata-Version: 2.3
+Name: nanodrr
+Version: 0.1.0
+Summary: Blazing fast differentiable DRR rendering in modern PyTorch
+Requires-Dist: jaxtyping>=0.3.0
+Requires-Dist: matplotlib>=3.0.0
+Requires-Dist: roma>=1.5.6
+Requires-Dist: torch>=2.4.0
+Requires-Dist: torchio>=0.21.0
+Requires-Dist: pyvista[all]>=0.47.0 ; extra == 'scene'
+Requires-Dist: vtk>=9.6.0 ; extra == 'scene'
+Requires-Python: >=3.10
+Provides-Extra: scene
+Description-Content-Type: text/markdown
+
+# nanodrr
+
+A performance-oriented reimplementation of [`DiffDRR`](https://github.com/eigenvivek/DiffDRR) with the following improvements:
+
+- Optimized, pure PyTorch implementation (**~5× faster than `DiffDRR` at baseline**)
+- Modular design (freely swap subjects, extrinsics, and intrinsics during rendering)
+- Compatibility with `torch.compile` and mixed precision
+- Extensive type hints with `jaxtyping`
+- Standard Python package structure managed with `uv`
+
+All projective geometry is implemented internally using the standard [Hartley and Zisserman](https://www.cambridge.org/core/books/multiple-view-geometry-in-computer-vision/0B6F289C78B2B23F596CAA76D3D43F7A) pinhole camera formulation.
+
+## Installation
+
+> [!NOTE]
+>
+> On `pytorch<2.9`, `torch.compile` with `bfloat16` is slower than eager due to a CUDA graph capture issue (see [Benchmarks](#benchmarks)). Use `pytorch>=2.9` (Triton ≥3.5) for best results.
+
+```
+pip install "git+https://github.com/eigenvivek/nanodrr.git"
+```
+
+## Benchmarks
+
+> [!IMPORTANT]
+> - **~5× faster** than [`DiffDRR`](https://github.com/eigenvivek/DiffDRR) out of the box, without compilation (946 FPS vs 213 FPS)
+> - **~8× faster** with `torch.compile` and `bfloat16` on `pytorch>=2.9` (1,650 FPS vs 213 FPS)
+> - **~2.5× less memory** than `DiffDRR` (516 MB vs 1,344 MB peak reserved with `bfloat16` + compile)
+
+![Benchmarking runtime, FPS, and memory usage.](tests/benchmark/benchmark.png "benchmark")
+
+> *Mean ± std. dev. of 10 runs, 100 loops each. Benchmarked by rendering 200×200 DRRs on an NVIDIA RTX 6000 Ada (48 GB) with Python 3.12. Compile represents `torch.compile(mode="reduce-overhead", fullgraph=True)`. Full experiment at [`tests/benchmark/`](tests/benchmark/).*
+
+## Roadmap
+
+- [x] Implement a fully optimized renderer
+- [x] Port strictly necessary modules from `DiffDRR` (e.g., SE(3) utilities, loss functions, and 2D plotting)
+- [x] Migrate 3D plotting functions to an optional module
+- [ ] Integrate with [`xvr`](https://github.com/eigenvivek/xvr) to speed up network training and registration
+- [ ] Integrate with [`polypose`](https://github.com/eigenvivek/polypose) to speed up registration
+- [ ] Release as `v1.0.0` of `DiffDRR`!

nanodrr-0.1.0/README.md

@@ -0,0 +1,41 @@
+# nanodrr
+
+A performance-oriented reimplementation of [`DiffDRR`](https://github.com/eigenvivek/DiffDRR) with the following improvements:
+
+- Optimized, pure PyTorch implementation (**~5× faster than `DiffDRR` at baseline**)
+- Modular design (freely swap subjects, extrinsics, and intrinsics during rendering)
+- Compatibility with `torch.compile` and mixed precision
+- Extensive type hints with `jaxtyping`
+- Standard Python package structure managed with `uv`
+
+All projective geometry is implemented internally using the standard [Hartley and Zisserman](https://www.cambridge.org/core/books/multiple-view-geometry-in-computer-vision/0B6F289C78B2B23F596CAA76D3D43F7A) pinhole camera formulation.
+
+## Installation
+
+> [!NOTE]
+>
+> On `pytorch<2.9`, `torch.compile` with `bfloat16` is slower than eager due to a CUDA graph capture issue (see [Benchmarks](#benchmarks)). Use `pytorch>=2.9` (Triton ≥3.5) for best results.
+
+```
+pip install "git+https://github.com/eigenvivek/nanodrr.git"
+```
+
+## Benchmarks
+
+> [!IMPORTANT]
+> - **~5× faster** than [`DiffDRR`](https://github.com/eigenvivek/DiffDRR) out of the box, without compilation (946 FPS vs 213 FPS)
+> - **~8× faster** with `torch.compile` and `bfloat16` on `pytorch>=2.9` (1,650 FPS vs 213 FPS)
+> - **~2.5× less memory** than `DiffDRR` (516 MB vs 1,344 MB peak reserved with `bfloat16` + compile)
+
+![Benchmarking runtime, FPS, and memory usage.](tests/benchmark/benchmark.png "benchmark")
+
+> *Mean ± std. dev. of 10 runs, 100 loops each. Benchmarked by rendering 200×200 DRRs on an NVIDIA RTX 6000 Ada (48 GB) with Python 3.12. Compile represents `torch.compile(mode="reduce-overhead", fullgraph=True)`. Full experiment at [`tests/benchmark/`](tests/benchmark/).*
+
+## Roadmap
+
+- [x] Implement a fully optimized renderer
+- [x] Port strictly necessary modules from `DiffDRR` (e.g., SE(3) utilities, loss functions, and 2D plotting)
+- [x] Migrate 3D plotting functions to an optional module
+- [ ] Integrate with [`xvr`](https://github.com/eigenvivek/xvr) to speed up network training and registration
+- [ ] Integrate with [`polypose`](https://github.com/eigenvivek/polypose) to speed up registration
+- [ ] Release as `v1.0.0` of `DiffDRR`!

nanodrr-0.1.0/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "nanodrr"
+version = "0.1.0"
+description = "Blazing fast differentiable DRR rendering in modern PyTorch"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "jaxtyping>=0.3.0",
+    "matplotlib>=3.0.0",
+    "roma>=1.5.6",
+    "torch>=2.4.0",
+    "torchio>=0.21.0",
+]
+
+[project.optional-dependencies]
+scene = [
+    "pyvista[all]>=0.47.0",
+    "vtk>=9.6.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.10.2,<0.11.0"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 120
+
+[tool.ruff.lint]
+ignore = [
+    "F722",  # Forward annotation false positive from jaxtyping
+    "F821",  # Forward annotation false positive from jaxtyping
+]
+
+[dependency-groups]
+dev = [
+    "prek>=0.3.2",
+]
+docs = [
+    "griffe>=2.0.0",
+    "mkdocs-callouts>=1.16.0",
+    "mkdocs-gen-files>=0.6.0",
+    "mkdocs-jupyter>=0.25.1",
+    "mkdocs-literate-nav>=0.6.2",
+    "mkdocs-material>=9.7.1",
+    "mkdocs-same-dir>=0.1.3",
+    "mkdocs-section-index>=0.3.10",
+    "mkdocstrings[python]>=1.0.3",
+]

nanodrr-0.1.0/src/nanodrr/camera/.ipynb_checkpoints/extrinsics-checkpoint.py

@@ -0,0 +1,70 @@
+import torch
+from jaxtyping import Float
+
+from ..geometry import convert
+
+_ORIENTATION_MATRICES = {
+    "AP": [
+        [-1, 0,  0, 0],
+        [ 0, 0, -1, 0],
+        [ 0,-1,  0, 0],
+        [ 0, 0,  0, 1],
+    ],
+    "PA": [
+        [-1, 0, 0, 0],
+        [ 0, 0, 1, 0],
+        [ 0,-1, 0, 0],
+        [ 0, 0, 0, 1],
+    ],
+    None: [
+        [-1, 0, 0, 0],
+        [ 0,-1, 0, 0],
+        [ 0, 0, 1, 0],
+        [ 0, 0, 0, 1],
+    ],
+}
+
+
+def make_rt_inv(
+    rotation: Float[torch.Tensor, "B 3"],
+    translation: Float[torch.Tensor, "B 3"],
+    orientation: str | None = "AP",
+    isocenter: Float[torch.Tensor, "3"] | None = None,
+) -> Float[torch.Tensor, "B 4 4"]:
+    """Create 4x4 camera-to-world (extrinsic inverse) matrices.
+
+    Composes pose and reorientation as ``extrinsic_inv = pose @ reorient``
+    so that *translation* is applied in the pre-reoriented frame.
+
+    Args:
+        rotation: (B, 3) Euler angles (z, x, y) in degrees, ZXY convention.
+        translation: (B, 3) camera position in mm, relative to *isocenter*
+                     (or world origin when isocenter is ``None``).
+        orientation: ``"AP"``, ``"PA"``, or ``None``.
+        isocenter: Optional (3,) volume centre in world coordinates.
+
+    Returns:
+        (B, 4, 4) camera-to-world transformation matrices.
+    """
+    if orientation not in _ORIENTATION_MATRICES:
+        raise ValueError(f"Unknown orientation: {orientation}. Use 'AP', 'PA', or None")
+
+    device = rotation.device
+    dtype = rotation.dtype
+
+    if isocenter is None:
+        isocenter = torch.zeros(3, device=device, dtype=dtype)
+
+    pose = convert(rotation, translation, "euler", convention="ZXY", isocenter=isocenter)
+    orientation_matrix = _get_orientation_matrix(orientation, device, dtype)
+
+    return pose @ orientation_matrix
+
+
+def _get_orientation_matrix(
+    orientation: str | None,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> Float[torch.Tensor, "4 4"]:
+    """Return the combined orientation + Rz(180°) matrix."""
+    return torch.tensor(_ORIENTATION_MATRICES[orientation], device=device, dtype=dtype)

nanodrr-0.1.0/src/nanodrr/camera/.ipynb_checkpoints/intrinsics-checkpoint.py

@@ -0,0 +1,55 @@
+import torch
+from jaxtyping import Float
+
+
+def make_k_inv(
+    sdd: float,
+    delx: float,
+    dely: float,
+    x0: float,
+    y0: float,
+    height: int,
+    width: int,
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
+) -> Float[torch.Tensor, "1 3 3"]:
+    """Build the inverse intrinsic matrix K⁻¹ for a cone-beam projector.
+
+    Focal lengths and principal point are derived from the physical geometry:
+
+        fx = sdd / delx          cy = y0 / dely + height / 2
+        fy = sdd / dely          cx = x0 / delx + width  / 2
+
+    The returned matrix is the analytical inverse of:
+
+        K = [[fx, 0, cx],
+             [0, fy, cy],
+             [0,  0,  1]]
+
+    Args:
+        sdd: Source-to-detector distance (mm).
+        delx, dely: Pixel spacing in x and y (mm/px).
+        x0, y0: Principal-point offset from detector centre (mm).
+        height, width: Detector dimensions in pixels.
+        dtype: Optional tensor dtype.
+        device: Optional tensor device.
+
+    Returns:
+        (1, 3, 3) inverse intrinsic matrix.
+    """
+    fx = sdd / delx
+    fy = sdd / dely
+    cx = x0 / delx + width / 2.0
+    cy = y0 / dely + height / 2.0
+
+    return torch.tensor(
+        [
+            [
+                [1.0 / fx, 0.0, -cx / fx],
+                [0.0, 1.0 / fy, -cy / fy],
+                [0.0, 0.0, 1.0],
+            ]
+        ],
+        dtype=dtype,
+        device=device,
+    )

nanodrr-0.1.0/src/nanodrr/camera/.ipynb_checkpoints/matrices-checkpoint.py

@@ -0,0 +1,174 @@
+import torch
+
+
+def make_k_inv(
+    sdd: float,
+    delx: float,
+    dely: float,
+    x0: float,
+    y0: float,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    fx = sdd / delx
+    fy = sdd / dely
+    cx = x0 / delx + width / 2.0
+    cy = y0 / dely + height / 2.0
+
+    fx_inv = 1.0 / fx
+    fy_inv = 1.0 / fy
+
+    return torch.tensor(
+        [
+            [
+                [fx_inv, 0.0, -cx * fx_inv],
+                [0.0, fy_inv, -cy * fy_inv],
+                [0.0, 0.0, 1.0],
+            ]
+        ]
+    )
+
+
+def make_rt_inv(
+    rotation: torch.Tensor,
+    translation: torch.Tensor,
+    orientation: str | None = "AP",
+    isocenter: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Create 4x4 camera-to-world (extrinsic inverse) transformation matrix.
+
+    Composes the pose and reorientation to match DiffDRR's behavior:
+        extrinsic_inv = pose @ reorient
+
+    This order means the translation is applied in the pre-reoriented frame,
+    so translation=(0, 850, 0) with AP orientation places the source at Y=850
+    in world coordinates (behind the patient for AP imaging).
+
+    When isocenter is provided, the translation is interpreted as relative to
+    the isocenter rather than world origin.
+
+    Args:
+        rotation: (batch, 3) Euler angles (angle_z, angle_x, angle_y) in degrees, ZXY convention
+        translation: (batch, 3) camera position (mm). If isocenter is provided,
+                    this is relative to isocenter; otherwise relative to world origin.
+        orientation: "AP", "PA", or None for frame-of-reference
+        isocenter: Optional (3,) volume isocenter in world coordinates.
+                  When provided, the translation is relative to this point.
+
+    Returns:
+        (batch, 4, 4) camera-to-world transformation matrices
+    """
+    if orientation not in (None, "AP", "PA"):
+        raise ValueError(f"Unknown orientation: {orientation}. Use 'AP', 'PA', or None")
+
+    batch_size = rotation.shape[0]
+    device = rotation.device
+    dtype = rotation.dtype
+
+    # Default isocenter to origin
+    if isocenter is None:
+        isocenter = torch.zeros(3, device=device, dtype=dtype)
+
+    # Get rotation matrices from Euler angles
+    R = euler_to_matrix(rotation)  # (batch, 3, 3)
+
+    # Compute camera center: R @ translation + isocenter
+    # bij,bj->bi : batched matrix-vector multiply
+    camera_center = torch.einsum("bij,bj->bi", R, translation)
+    camera_center = camera_center + isocenter
+
+    # Build 4x4 pose matrices [R | camera_center]
+    pose = torch.zeros(batch_size, 4, 4, device=device, dtype=dtype)
+    pose[:, :3, :3] = R
+    pose[:, :3, 3] = camera_center
+    pose[:, 3, 3] = 1.0
+
+    # Apply orientation (pose @ combined)
+    # bij,jk->bik : batched matrix times single matrix
+    orientation_matrix = get_orientation_matrix(orientation, device, dtype)
+    out = torch.einsum("bij,jk->bik", pose, orientation_matrix)
+
+    return out
+
+
+def euler_to_matrix(rotation: torch.Tensor) -> torch.Tensor:
+    """Convert ZXY Euler angles (degrees) to rotation matrices.
+
+    Args:
+        rotation: Euler angles (angle_z, angle_x, angle_y) in degrees, shape (batch, 3)
+
+    Returns:
+        Rotation matrices of shape (batch, 3, 3)
+    """
+    angles = torch.deg2rad(rotation)
+    z, x, y = angles[:, 0], angles[:, 1], angles[:, 2]
+
+    cz, sz = torch.cos(z), torch.sin(z)
+    cx, sx = torch.cos(x), torch.sin(x)
+    cy, sy = torch.cos(y), torch.sin(y)
+
+    # ZXY Euler rotation matrix
+    R = torch.stack(
+        [
+            torch.stack(
+                [cy * cz - sx * sy * sz, -cx * sz, cz * sy + cy * sx * sz], dim=1
+            ),
+            torch.stack(
+                [cy * sz + cz * sx * sy, cx * cz, sy * sz - cy * cz * sx], dim=1
+            ),
+            torch.stack([-cx * sy, sx, cx * cy], dim=1),
+        ],
+        dim=1,
+    )
+
+    return R
+
+
+def get_orientation_matrix(
+    orientation: str | None, device: torch.device, dtype: torch.dtype
+) -> torch.Tensor:
+    """Get the combined orientation + Rz(180°) matrix.
+
+    Args:
+        orientation: "AP", "PA", or None
+        device: torch device
+        dtype: torch dtype
+
+    Returns:
+        4x4 transformation matrix
+    """
+    if orientation == "AP":
+        combined = torch.tensor(
+            [
+                [-1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, -1.0, 0.0],
+                [0.0, -1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+    elif orientation == "PA":
+        combined = torch.tensor(
+            [
+                [-1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, -1.0, 0.0],
+                [0.0, -1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+    else:  # None - just Rz180
+        combined = torch.tensor(
+            [
+                [-1.0, 0.0, 0.0, 0.0],
+                [0.0, -1.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+            device=device,
+            dtype=dtype,
+        )
+
+    return combined

nanodrr-0.1.0/src/nanodrr/camera/__init__.py

@@ -0,0 +1,4 @@
+from .intrinsics import make_k_inv
+from .extrinsics import make_rt_inv
+
+__all__ = ["make_k_inv", "make_rt_inv"]

nanodrr-0.1.0/src/nanodrr/camera/extrinsics.py

@@ -0,0 +1,70 @@
+import torch
+from jaxtyping import Float
+
+from ..geometry import convert
+
+_ORIENTATION_MATRICES = {
+    "AP": [
+        [-1, 0, 0, 0],
+        [0, 0, -1, 0],
+        [0, -1, 0, 0],
+        [0, 0, 0, 1],
+    ],
+    "PA": [
+        [-1, 0, 0, 0],
+        [0, 0, 1, 0],
+        [0, -1, 0, 0],
+        [0, 0, 0, 1],
+    ],
+    None: [
+        [-1, 0, 0, 0],
+        [0, -1, 0, 0],
+        [0, 0, 1, 0],
+        [0, 0, 0, 1],
+    ],
+}
+
+
+def make_rt_inv(
+    rotation: Float[torch.Tensor, "B 3"],
+    translation: Float[torch.Tensor, "B 3"],
+    orientation: str | None = "AP",
+    isocenter: Float[torch.Tensor, "3"] | None = None,
+) -> Float[torch.Tensor, "B 4 4"]:
+    """Create 4x4 camera-to-world (extrinsic inverse) matrices.
+
+    Composes pose and reorientation as ``extrinsic_inv = pose @ reorient``
+    so that *translation* is applied in the pre-reoriented frame.
+
+    Args:
+        rotation: (B, 3) Euler angles (z, x, y) in degrees, ZXY convention.
+        translation: (B, 3) camera position in mm, relative to *isocenter*
+                     (or world origin when isocenter is ``None``).
+        orientation: ``"AP"``, ``"PA"``, or ``None``.
+        isocenter: Optional (3,) volume centre in world coordinates.
+
+    Returns:
+        (B, 4, 4) camera-to-world transformation matrices.
+    """
+    if orientation not in _ORIENTATION_MATRICES:
+        raise ValueError(f"Unknown orientation: {orientation}. Use 'AP', 'PA', or None")
+
+    device = rotation.device
+    dtype = rotation.dtype
+
+    if isocenter is None:
+        isocenter = torch.zeros(3, device=device, dtype=dtype)
+
+    pose = convert(rotation, translation, "euler", convention="ZXY", isocenter=isocenter)
+    orientation_matrix = _get_orientation_matrix(orientation, device, dtype)
+
+    return pose @ orientation_matrix
+
+
+def _get_orientation_matrix(
+    orientation: str | None,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> Float[torch.Tensor, "4 4"]:
+    """Return the combined orientation + Rz(180°) matrix."""
+    return torch.tensor(_ORIENTATION_MATRICES[orientation], device=device, dtype=dtype)

nanodrr-0.1.0/src/nanodrr/camera/intrinsics.py

@@ -0,0 +1,57 @@
+import torch
+from jaxtyping import Float
+
+
+def make_k_inv(
+    sdd: float,
+    delx: float,
+    dely: float,
+    x0: float,
+    y0: float,
+    height: int,
+    width: int,
+    dtype: torch.dtype | None = None,
+    device: torch.device | None = None,
+) -> Float[torch.Tensor, "1 3 3"]:
+    """Build the inverse intrinsic matrix K⁻¹ for a cone-beam projector.
+
+    Focal lengths and principal point are derived from the physical geometry:
+
+        fx = sdd / delx          cy = y0 / dely + height / 2
+        fy = sdd / dely          cx = x0 / delx + width  / 2
+
+    The returned matrix is the analytical inverse of:
+
+        K = [[fx, 0, cx],
+             [0, fy, cy],
+             [0,  0,  1]]
+
+    Args:
+        sdd: Source-to-detector distance (mm).
+        delx: Pixel spacing in x (mm/px).
+        dely: Pixel spacing in y (mm/px).
+        x0: Principal-point offset from detector centre in x (mm).
+        y0: Principal-point offset from detector centre in y (mm).
+        height: Detector height in pixels.
+        width: Detector width in pixels.
+        dtype: Optional tensor dtype.
+        device: Optional tensor device.
+    Returns:
+        (1, 3, 3) inverse intrinsic matrix.
+    """
+    fx = sdd / delx
+    fy = sdd / dely
+    cx = x0 / delx + width / 2.0
+    cy = y0 / dely + height / 2.0
+
+    return torch.tensor(
+        [
+            [
+                [1.0 / fx, 0.0, -cx / fx],
+                [0.0, 1.0 / fy, -cy / fy],
+                [0.0, 0.0, 1.0],
+            ]
+        ],
+        dtype=dtype,
+        device=device,
+    )

nanodrr-0.1.0/src/nanodrr/data/.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1,4 @@
+from .demo import download_deepfluoro
+from .io import Subject
+
+__all__ = ["download_deepfluoro", "Subject"]

nanodrr-0.1.0/src/nanodrr/data/.ipynb_checkpoints/demo-checkpoint.py

@@ -0,0 +1,22 @@
+import os
+import torch
+from platformdirs import user_cache_dir
+
+CACHE_DIR = user_cache_dir("nanodrr")
+
+def download_deepfluoro(subject: int = 1) -> tuple[str, str]:
+    """Download a subject from the DeepFluoro dataset."""
+    subject = f"subject{subject:02d}"
+    base_url = f"https://huggingface.co/datasets/eigenvivek/xvr-data/resolve/main/deepfluoro/{subject}"
+    imagepath = os.path.join(CACHE_DIR, "deepfluoro", subject, "volume.nii.gz")
+    labelpath = os.path.join(CACHE_DIR, "deepfluoro", subject, "mask.nii.gz")
+
+    for url, local_path in [
+        (f"{base_url}/volume.nii.gz", imagepath),
+        (f"{base_url}/mask.nii.gz", labelpath),
+    ]:
+        if not os.path.exists(local_path):
+            os.makedirs(os.path.dirname(local_path), exist_ok=True)
+            torch.hub.download_url_to_file(url, local_path)
+
+    return imagepath, labelpath