Merge pull request #139 from eigenvivek/diffrend

Differentiable rendering and large DRRs

Author: eigenvivek

README.md

@@ -1,49 +1,39 @@
 DiffDRR
 ================
 
-<!-- WARNING: THIS FILE WAS AUTOGENERATED! DO NOT EDIT! -->
-
 > Auto-differentiable DRR synthesis and optimization in PyTorch
 
 [![CI](https://github.com/eigenvivek/DiffDRR/actions/workflows/test.yaml/badge.svg)](https://github.com/eigenvivek/DiffDRR/actions/workflows/test.yaml)
-[![Paper
-shield](https://img.shields.io/badge/arXiv-2208.12737-red.svg)](https://arxiv.org/abs/2208.12737)
-[![License:
-MIT](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
-[![Downloads](https://static.pepy.tech/personalized-badge/diffdrr?period=month&units=international_system&left_color=grey&right_color=blue&left_text=downloads.month)](https://pepy.tech/project/diffdrr)
+[![Paper shield](https://img.shields.io/badge/arXiv-2208.12737-red.svg)](https://arxiv.org/abs/2208.12737)
+[![License: MIT](https://img.shields.io/badge/License-Apache2.0-blue.svg)](LICENSE)
+[![Downloads](https://static.pepy.tech/personalized-badge/diffdrr?period=total&units=none&left_color=grey&right_color=blue&left_text=downloads)](https://pepy.tech/project/diffdrr)
 [![Docs](https://img.shields.io/badge/docs-passing-brightgreen.svg)](https://vivekg.dev/DiffDRR/)
-[![Code style:
-black](https://img.shields.io/badge/Code%20style-black-black.svg)](https://github.com/psf/black)
+[![Code style: black](https://img.shields.io/badge/Code%20style-black-black.svg)](https://github.com/psf/black)
 
-`DiffDRR` is a PyTorch-based digitally reconstructed radiograph (DRR)
-generator that provides
+`DiffDRR` is a PyTorch-based digitally reconstructed radiograph (DRR) generator that provides
 
-1.  Auto-differentiable DRR syntheisis
-2.  GPU-accelerated rendering
-3.  A pure Python implementation
+1. Auto-differentiable DRR syntheisis
+2. GPU-accelerated rendering
+3. A pure Python implementation
 
-Most importantly, `DiffDRR` implements DRR synthesis as a PyTorch
-module, making it interoperable in deep learning pipelines.
+Most importantly, `DiffDRR` implements DRR synthesis as a PyTorch module, making it interoperable in deep learning pipelines.
 
 - [Installation Guide](#installation-guide)
 - [Usage](#usage)
-- [Example: Rigid 2D-to-3D
-  registration](#application-6-dof-slice-to-volume-registration)
+- [Example: Rigid 2D-to-3D registration](#application-6-dof-slice-to-volume-registration)
 - [How does `DiffDRR` work?](#how-does-diffdrr-work)
 - [Citing `DiffDRR`](#citing-diffdrr)
 
 ## Installation Guide
 
 To install `DiffDRR` from PyPI:
-
-``` zsh
+```zsh
 pip install diffdrr
 ```
 
 ## Usage
 
-The following minimal example specifies the geometry of the projectional
-radiograph imaging system and traces rays through a CT volume:
+The following minimal example specifies the geometry of the projectional radiograph imaging system and traces rays through a CT volume:
 
 ``` python
 import matplotlib.pyplot as plt
@@ -57,43 +47,34 @@ from diffdrr.visualization import plot_drr
 volume, spacing = load_example_ct()
 
 # Initialize the DRR module for generating synthetic X-rays
+device = "cuda" if torch.cuda.is_available() else "cpu"
 drr = DRR(
-    volume,        # The CT volume as a numpy array
-    spacing,       # Voxel dimensions of the CT
-    sdr=300.0,     # Source-to-detector radius (half of the source-to-detector distance)
-    height=200,    # Height of the DRR (if width is not seperately provided, the generated image is square)
-    delx=4.0,      # Pixel spacing (in mm)
-    batch_size=1,  # How many batches of parameters will be passed = number of DRRs generated each forward pass
-).to("cuda" if torch.cuda.is_available() else "cpu")
-
-# Rotations and translations determine the viewing angle
-# They must have the same batch_size as was passed to the DRR constructor
-# Rotations are (yaw pitch roll)
-# Translations are (bx by bz)
-rotations = torch.tensor([[torch.pi, 0.0, torch.pi / 2]])
-translations = torch.tensor(volume.shape) * torch.tensor(spacing) / 2
-translations = translations.unsqueeze(0)
-
-# Generate the DRR
-drr.move_carm(rotations, translations)
-with torch.no_grad():
-    img = drr()  # Only keep the graph if optimizing DRRs
-ax = plot_drr(img)
+    volume,      # The CT volume as a numpy array
+    spacing,     # Voxel dimensions of the CT
+    sdr=300.0,   # Source-to-detector radius (half of the source-to-detector distance)
+    height=200,  # Height of the DRR (if width is not seperately provided, the generated image is square)
+    delx=4.0,    # Pixel spacing (in mm)
+).to(device)
+
+# Set the camera pose with rotations (yaw, pitch, roll) and translations (x, y, z)
+rotations = torch.tensor([[torch.pi, 0.0, torch.pi / 2]], device=device)
+bx, by, bz = torch.tensor(volume.shape) * torch.tensor(spacing) / 2
+translations = torch.tensor([[bx, by, bz]], device=device)
+
+# Make the DRR
+img = drr(rotations, translations)
+plot_drr(img, ticks=False)
 plt.show()
 ```
 
-![](index_files/figure-commonmark/cell-2-output-1.png)
+![](notebooks/index_files/figure-commonmark/cell-2-output-1.png)
 
 On a single NVIDIA RTX 2080 Ti GPU, producing such an image takes
 
-``` python
-%timeit drr()
-```
-
-    34.9 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
+    34.9 ms ± 32.2 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
 
 The full example is available at
-[`tutorials/introduction.ipynb`](tutorials/introduction.ipynb).
+[`introduction.ipynb`](https://vivekg.dev/DiffDRR/tutorials/introduction.html).
 
 ## Application: 6-DoF Slice-to-Volume Registration
 
@@ -111,7 +92,7 @@ optimization runs like this:
 ![](https://cdn.githubraw.com/eigenvivek/DiffDRR/7a6a44aeab58d19cc7a4afabfc5aabab3a494974/experiments/registration/results/momentum_dampen/gifs/converged/649.gif)
 
 The full example is available at
-[`tutorials/optimizers.ipynb`](tutorials/optimizers.ipynb).
+[`optimizers.ipynb`](https://vivekg.dev/DiffDRR/tutorials/optimizers.html).
 
 ## How does `DiffDRR` work?
 

diffdrr/__init__.py

@@ -1 +1 @@
-__version__ = "0.3.2"
+__version__ = "0.3.3"

diffdrr/_modidx.py

@@ -18,8 +18,10 @@
             'diffdrr.drr': { 'diffdrr.drr.DRR': ('api/drr.html#drr', 'diffdrr/drr.py'),
                              'diffdrr.drr.DRR.__init__': ('api/drr.html#drr.__init__', 'diffdrr/drr.py'),
                              'diffdrr.drr.DRR.forward': ('api/drr.html#drr.forward', 'diffdrr/drr.py'),
-                             'diffdrr.drr.DRR.move_carm': ('api/drr.html#drr.move_carm', 'diffdrr/drr.py'),
-                             'diffdrr.drr.DRR.reshape_transform': ('api/drr.html#drr.reshape_transform', 'diffdrr/drr.py')},
+                             'diffdrr.drr.DRR.reshape_transform': ('api/drr.html#drr.reshape_transform', 'diffdrr/drr.py'),
+                             'diffdrr.drr.Registration': ('api/drr.html#registration', 'diffdrr/drr.py'),
+                             'diffdrr.drr.Registration.__init__': ('api/drr.html#registration.__init__', 'diffdrr/drr.py'),
+                             'diffdrr.drr.Registration.forward': ('api/drr.html#registration.forward', 'diffdrr/drr.py')},
             'diffdrr.metrics': { 'diffdrr.metrics.XCorr2': ('api/metrics.html#xcorr2', 'diffdrr/metrics.py'),
                                  'diffdrr.metrics.XCorr2.__init__': ('api/metrics.html#xcorr2.__init__', 'diffdrr/metrics.py'),
                                  'diffdrr.metrics.XCorr2.forward': ('api/metrics.html#xcorr2.forward', 'diffdrr/metrics.py')},

diffdrr/drr.py

@@ -6,19 +6,18 @@
 import numpy as np
 import torch
 import torch.nn as nn
-
 from fastcore.basics import patch
 
-from .siddon import siddon_raycast
 from .detector import Detector
+from .siddon import siddon_raycast
 from .utils import reshape_subsampled_drr
 
 # %% auto 0
-__all__ = ['DRR']
+__all__ = ['DRR', 'Registration']
 
 # %% ../notebooks/api/00_drr.ipynb 5
 class DRR(nn.Module):
-    """Torch module that computes differentiable digitally reconstructed radiographs."""
+    """PyTorch module that computes differentiable digitally reconstructed radiographs."""
 
     def __init__(
         self,
@@ -31,16 +30,13 @@ def __init__(
         | None = None,  # Width of the rendered DRR (if not provided, set to `height`)
         dely: float | None = None,  # Y-axis pixel size (if not provided, set to `delx`)
         p_subsample: float | None = None,  # Proportion of pixels to randomly subsample
-        reshape: bool = True,  # Return DRR with shape (b, h, w)
+        reshape: bool = True,  # Return DRR with shape (b, 1, h, w)
         convention: str = "diffdrr",  # Either `diffdrr` or `deepdrr`, order of basis matrix multiplication
-        batch_size: int = 1,  # Number of DRRs to generate per forward pass
+        patch_size: int
+        | None = None,  # Render patches of the DRR in series (useful for large DRRs)
     ):
         super().__init__()
 
-        params = torch.empty(batch_size, 6)
-        self.rotations = nn.Parameter(params[..., :3])
-        self.translations = nn.Parameter(params[..., 3:])
-
         # Initialize the X-ray detector
         width = height if width is None else width
         dely = delx if dely is None else dely
@@ -60,9 +56,9 @@ def __init__(
         self.register_buffer("spacing", torch.tensor(spacing))
         self.register_buffer("volume", torch.tensor(volume).flip([0]))
         self.reshape = reshape
-
-        # Dummy tensor for device and dtype
-        self.register_buffer("dummy", torch.tensor([0.0]))
+        self.patch_size = patch_size
+        if self.patch_size is not None:
+            self.n_patches = (height * width) // (self.patch_size**2)
 
     def reshape_transform(self, img, batch_size):
         if self.reshape:
@@ -74,18 +70,41 @@ def reshape_transform(self, img, batch_size):
 
 # %% ../notebooks/api/00_drr.ipynb 7
 @patch
-def move_carm(self: DRR, rotations: torch.Tensor, translations: torch.Tensor):
-    state_dict = self.state_dict()
-    state_dict["rotations"].copy_(rotations)
-    state_dict["translations"].copy_(translations)
-
-# %% ../notebooks/api/00_drr.ipynb 8
-@patch
-def forward(self: DRR):
+def forward(self: DRR, rotations: torch.Tensor, translations: torch.Tensor):
     """Generate DRR with rotations and translations parameters."""
+    assert len(rotations) == len(translations)
+    batch_size = len(rotations)
     source, target = self.detector.make_xrays(
-        rotations=self.rotations,
-        translations=self.translations,
+        rotations=rotations,
+        translations=translations,
     )
-    img = siddon_raycast(source, target, self.volume, self.spacing)
-    return self.reshape_transform(img, batch_size=len(self.rotations))
+
+    if self.patch_size is not None:
+        n_points = target.shape[1] // self.n_patches
+        img = []
+        for idx in range(self.n_patches):
+            t = target[:, idx * n_points : (idx + 1) * n_points]
+            partial = siddon_raycast(source, t, self.volume, self.spacing)
+            img.append(partial)
+        img = torch.cat(img, dim=1)
+    else:
+        img = siddon_raycast(source, target, self.volume, self.spacing)
+    return self.reshape_transform(img, batch_size=batch_size)
+
+# %% ../notebooks/api/00_drr.ipynb 9
+class Registration(nn.Module):
+    """Perform automatic 2D-to-3D registration using differentiable rendering."""
+
+    def __init__(
+        self,
+        drr: DRR,
+        rotations: torch.Tensor,
+        translations: torch.Tensor,
+    ):
+        super().__init__()
+        self.drr = drr
+        self.rotations = nn.Parameter(rotations)
+        self.translations = nn.Parameter(translations)
+
+    def forward(self):
+        return self.drr(self.rotations, self.translations)

diffdrr/visualization.py

@@ -60,6 +60,7 @@ def animate(
     drr: DRR,
     ground_truth: torch.Tensor | None = None,
     verbose: bool = True,
+    device="cpu",
     **kwargs,  # To pass to imageio.v3.imwrite
 ):
     """Animate the optimization of a DRR."""
@@ -96,10 +97,15 @@ def make_fig(ground_truth):
         for idx, row in itr:
             fig, ax_opt = make_fig() if ground_truth is None else make_fig(ground_truth)
             params = row[["theta", "phi", "gamma", "bx", "by", "bz"]].values
-            rotations = torch.tensor(row[["theta", "phi", "gamma"]].values)
-            translations = torch.tensor(row[["bx", "by", "bz"]].values)
-            drr.move_carm(rotations, translations)
-            itr = drr().detach()
+            rotations = (
+                torch.tensor(row[["theta", "phi", "gamma"]].values)
+                .unsqueeze(0)
+                .to(device)
+            )
+            translations = (
+                torch.tensor(row[["bx", "by", "bz"]].values).unsqueeze(0).to(device)
+            )
+            itr = drr(rotations, translations)
             _ = plot_drr(itr, axs=ax_opt)
             ax_opt.set(xlabel="Moving DRR")
             fig.savefig(f"{tmpdir}/{idx}.png")

notebooks/_quarto.yml

@@ -11,9 +11,13 @@ website:
   twitter-card: true
   open-graph: true
   repo-actions: [issue]
+  favicon: favicon.png
   navbar:
     background: primary
     search: true
+    right:
+      - icon: github
+        href: "https://github.com/eigenvivek/DiffDRR"
   sidebar:
     style: floating