repository·main·Indexed 42 days ago
https://github.com/eigenvivek/diffdrrA PyTorch-based, auto-differentiable digitally reconstructed radiograph (DRR) generator. It enables GPU-accelerated, physics-based X-ray synthesis and optimization for deep learning pipelines involving 2D/3D registration, segmentation, and volume reconstruction. The library supports multiple SO(3) rotation representations, including Quaternions, Rotation matrices, Axis-angle, and Euler angles.
Both Siddon and Trilinear renderers support an optional mask argument in their forward methods.
When a mask is provided, the renderer uses the mask to perform a scatter_add_ operation, effectively grouping the sampled voxel values into different channels based on the mask's values. This is useful for segmenting or isolating specific structures within the volume during the rendering process.
The DRR module supports multiple representations for the rotation component of the camera pose. When calling the module, use the parameterization argument to select the format and convention to specify the axis ordering (e.g., for Euler angles). Supported formats include:
parameterization="euler_angles", convention="ZXY")To generate synthetic X-rays, use the DRR class from diffdrr.drr. You must provide a subject (containing the CT volume, origin, and voxel spacing) and specify imaging geometry parameters like sdd (source-to-detector distance) and pixel spacing (delx).
The DRR module is a PyTorch module, so you can move it to a GPU using .to(device). When calling the module, you provide rotations and translations. You can specify the rotation representation using the parameterization and convention arguments.
import torch
from diffdrr.drr import DRR
from diffdrr.data import load_example_ct
# 1. Load subject data
subject = load_example_ct()
# 2. Initialize DRR module
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
drr = DRR(
subject, # Object with CT volume, origin, and spacing
sdd=1020.0, # Source-to-detector distance
height=200, # Image height
delx=2.0, # Pixel spacing (mm)
).to(device)
# 3. Define camera pose
# Rotations can be euler angles, quaternions, rotation matrices, etc.
rotations = torch.tensor([[0.0, 0.0, 0.0]], device=device)
translations = torch.tensor([[0.0, 850.0, 0.0]], device=device)
# 4. Render
img = drr(rotations, translations, parameterization="euler_angles", convention="ZXY")You can install the latest stable release of diffdrr using pip. If you need the development version, clone the repository and install it in editable mode with the [dev] extra.
Note: For production use, the authors recommend nanodrr, which is a highly optimized reimplementation that is ~5x faster and uses 2.5x less memory while remaining pure PyTorch.
# Recommended: Install stable release
pip install diffdrr
# Development version
git clone https://github.com/eigenvivek/DiffDRR.git --depth 1
pip install -e 'DiffDRR/[dev]'DiffDRR can project 3D labelmaps into 2D using perspective geometry. This allows you to identify specific anatomical structures within simulated X-rays.
Because DiffDRR is differentiable with respect to both the camera poses and the 3D volume itself, it can be used for differentiable volume reconstruction (e.g., cone-beam CT reconstruction).
Both Siddon and Trilinear renderers support an optional mask argument in their forward methods. This allows you to render specific segments or channels of a volume separately.
When a mask is provided:
grid_sample with mode="nearest" to determine which channel/segment each sampled point belongs to.scatter_add_ to accumulate the sampled intensities into a tensor of shape (B, C, D), where C is the number of channels (determined by mask.max() + 1) and D is the number of rays.This is useful for multi-label segmentation tasks where you want to compute the X-ray projection of specific anatomical structures independently.
DiffDRR converts Hounsfield Units (HU) to a density-like quantity for rendering using a specific thresholding convention:
bone_attenuation_multiplier.Warning for Thoracic CT: Because aerated lung parenchyma typically spans $-950$ to $-700$ HU, much of the lung structure may be mapped to the same value as air (the $-800$ HU floor), making low-attenuation findings like emphysema indistinguishable. For applications requiring lung attenuation preservation, use nanodrr instead.
DiffDRR models X-ray imaging systems as pinhole cameras using the following parameters:
These define the camera's internal geometry and are often found in DICOM files:
sdd: Source-to-detector distance (focal length).delx: X-direction pixel spacing (mm).dely: Y-direction pixel spacing (mm).x0: Principal point offset in the x-direction.y0: Principal point offset in the y-direction.Extrinsic parameters define the camera's position and orientation in space as an affine transform $\mathbf T \in \mathbf{SE}(3)$. This is commonly parameterized using:
xyz)For implementation details on parameterization, refer to the diffdrr.pose module.
You can visualize the 3D geometry of the X-ray detector, camera pose, and CT volume using PyVista. This is useful for debugging the spatial relationship between the source, detector, and subject.
To enable visualization, ensure you have pyvista, trame, and vtk installed. The visualization process typically involves:
drr_to_mesh or labelmap_to_mesh.img_to_mesh.pyvista.Plotter.import pyvista
import torch
from diffdrr.data import load_example_ct
from diffdrr.drr import DRR
from diffdrr.pose import convert
from diffdrr.visualization import drr_to_mesh, img_to_mesh, labelmap_to_mesh, plot_drr
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 1. Load subject and create mesh
subject = load_example_ct()
ct_mesh = drr_to_mesh(subject, "surface_nets", threshold=225)
# 2. Initialize DRR and create pose
drr = DRR(subject, sdd=1020.0, height=200, delx=2.0).to(device)
rot = torch.tensor([[45.0, 30.0, 0.0]], device=device) / 180 * torch.pi
xyz = torch.tensor([[0.0, 800.0, 0.0]], device=device)
pose = convert(rot, xyz, parameterization="euler_angles", convention="ZXY")
# 3. Get geometry for plotting
camera, detector, texture, principal_ray = img_to_mesh(drr, pose)
# 4. Plot
plotter = pyvista.Plotter()
plotter.add_mesh(ct_mesh)
plotter.add_mesh(camera, show_edges=True)
plotter.add_mesh(detector, texture=texture)
plotter.show()The reducefn parameter in both Siddon and Trilinear renderers determines how the sampled values along a single ray are combined into a single pixel value.
Supported values:
"sum": Sums the contributions along the ray (standard for X-ray integration)."max": Takes the maximum value along the ray.Callable: You can pass a custom Python callable to define your own reduction logic.Initialize the Siddon module with the following options:
| Parameter | Type | Default | Description |
|---|---|---|---|
voxel_shift | float | 0.5 | 0 or 0.5, depending if the voxel is at the top left corner or the center |
mode | str | ` |