Skip to main content

Auto-differentiable digitally reconstructed radiographs in PyTorch

Project description

DiffDRR

Auto-differentiable DRR synthesis and optimization in PyTorch

CI Paper shield License: MIT Downloads Docs Code style: black

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

Most importantly, DiffDRR implements DRR synthesis as a PyTorch module, making it interoperable in deep learning pipelines.

Installation Guide

To install DiffDRR from PyPI:

pip install diffdrr

Development (optional)

DiffDRR source code, docs, and CI are all built using nbdev. To get set up with nbdev, install the following

mamba install jupyterlab nbdev -c fastai -c conda-forge 
nbdev_install_quarto  # To build docs
nbdev_install_hooks  # Make notebooks git-friendly

Running nbdev_help will give you the full list of options. The most important ones are

nbdev_preview  # Render docs locally and inspect in browser
nbdev_prepare  # NECESSARY BEFORE PUSHING: builds package, tests notebooks, and builds docs in one step

For more details, follow this in-depth tutorial.

Usage

The following minimal example specifies the geometry of the projectional radiograph imaging system and traces rays through a CT volume:

import matplotlib.pyplot as plt
import torch

from diffdrr.drr import DRR
from diffdrr.data import load_example_ct
from diffdrr.visualization import plot_drr

# Read in the volume and get the isocenter
volume, spacing = load_example_ct()
bx, by, bz = torch.tensor(volume.shape) * torch.tensor(spacing) / 2

# 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)
).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)
translations = torch.tensor([[bx, by, bz]], device=device)

# Make the DRR
img = drr(rotations, translations)
plot_drr(img, ticks=False)
plt.show()

On a single NVIDIA RTX 2080 Ti GPU, producing such an image takes

34.9 ms ± 32.2 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

The full example is available at introduction.ipynb.

Application: 6-DoF Slice-to-Volume Registration

We demonstrate the utility of our auto-differentiable DRR generator by solving a 6-DoF registration problem with gradient-based optimization. Here, we generate two DRRs:

  1. A fixed DRR from a set of ground truth parameters
  2. A moving DRR from randomly initialized parameters

To solve the registration problem, we use gradient descent to maximize an image loss similarity metric between the two DRRs. This produces optimization runs like this:

The full example is available at optimizers.ipynb.

How does DiffDRR work?

DiffDRR reformulates Siddon’s method,[^1] the canonical algorithm for calculating the radiologic path of an X-ray through a volume, as a series of vectorized tensor operations. This version of the algorithm is easily implemented in tensor algebra libraries like PyTorch to achieve a fast auto-differentiable DRR generator.

[^1]: Siddon RL. Fast calculation of the exact radiological path for a three-dimensional CT array. Medical Physics, 2(12):252–5, 1985.

Citing DiffDRR

If you find DiffDRR useful in your work, please cite our paper (or the freely accessible arXiv version):

@inproceedings{gopalakrishnanDiffDRR2022,
    author    = {Gopalakrishnan, Vivek and Golland, Polina},
    title     = {Fast Auto-Differentiable Digitally Reconstructed Radiographs for Solving Inverse Problems in Intraoperative Imaging},
    year      = {2022},
    booktitle = {Clinical Image-based Procedures: 11th International Workshop, CLIP 2022, Held in Conjunction with MICCAI 2022, Singapore, Proceedings},
    series    = {Lecture Notes in Computer Science},
    publisher = {Springer},
    doi       = {https://doi.org/10.1007/978-3-031-23179-7_1},
}

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

diffdrr-0.3.5.tar.gz (34.6 MB view hashes)

Uploaded Source

Built Distribution

diffdrr-0.3.5-py3-none-any.whl (34.6 MB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page