flashdrr.rendering

Differentiable DRR raycasting and volume rendering.

Raycasters

The package ships two raycaster back-ends that share the same camera model and VTK/RAS conventions:

  • :class:VolumeRaycaster — the main, fully differentiable raycaster. It integrates attenuation along each ray (Beer–Lambert by default) and supports batching over both volumes and views, optional memory checkpointing and tiled rendering, an optional learned scatter module, and an optional Poisson noise model.

  • :class:FusedVolumeRenderer — a Triton-fused CUDA kernel with a matching autograd wrapper. It is used internally by :class:VolumeRaycaster when triton=True and can also be used standalone for very fast projection.

VolumeRaycaster

class flashdrr.rendering.VolumeRaycaster[source]

Bases: Module

Differentiable C-arm volume raycaster for digitally reconstructed radiographs (DRRs).

Given a 3D attenuation volume in IJK space and one or more camera view matrices in RAS space, VolumeRaycaster samples along each camera ray, integrates the attenuation using either Beer-Lambert or alpha compositing, and produces a 2D projection per view. Rendering is differentiable with respect to the input volume and view matrices.

Two backends are supported:

  • A tiled Python-loop renderer using torch.grid_sample (default).

  • A fused CUDA/Triton renderer (flashdrr.rendering.FusedVolumeRenderer), selected per-call via forward(..., triton=True).

A lightweight, prototype depth-aware scatter model (DepthAwareScatter) can be attached for the last scatter channels of the volume.

__init__(density_factor=100.0, ray_samples=256, resolution=(224, 224), use_checkpointing=False, use_beer_lambert=True, scatter=None, i0=None, fov=20.0)[source]

Initialize the raycaster.

Parameters:

  • density_factor (float) – Scalar multiplier applied to volume intensities before integration. Used by the alpha-compositing branch of the Python-loop backend and passed to the Triton backend.

  • ray_samples (int) – Number of samples taken along each camera ray.

  • resolution (Tuple[int, int]) – Output image size as (width, height) or a single int for a square image.

  • use_checkpointing (bool) – If True, use torch.utils.checkpoint to trade compute for memory in the Python-loop backend (training only).

  • use_beer_lambert (bool) – If True, integrate via the Beer-Lambert law and return 1 - exp(-int(mu * dz)). If False, use front-to-back alpha compositing (over operator).

  • scatter (int | None) – If not None, treat the last scatter channels of the volume as scatter channels and process them with DepthAwareScatter (prototype; not fully functional); the remaining channels are rendered with Beer-Lambert as primary. Only has an effect when use_beer_lambert is True.

  • i0 (float | None) – Mean photon count at the source for the Gaussian-Poisson surrogate noise model. If None, no noise is added.

  • fov (float | None) – Vertical field of view in degrees used to build the per-pixel ray directions.

Return type:

None

apply_poisson(transmission)[source]

Add a differentiable Gaussian approximation of Poisson noise to a transmission map.

If self.i0 is None, the input is returned unchanged (no noise is added). Otherwise, Gaussian noise with standard deviation sqrt(transmission / i0) is added and the result is clamped to [0, 1]. This is a smooth surrogate for shot/Poisson noise and is fully differentiable.

Parameters:

transmission (torch.Tensor) – Transmission values (post-primary, pre-noise) in [0, 1]. Any shape is supported.

Returns:

Noisy transmission, same shape as transmission, clamped to [0, 1]. Returns the input unchanged when i0 is None.

Return type:

torch.Tensor

static compute_clipping_distances(camera_matrix, volume_shape, ijk2ras, margin=30.0)[source]

Compute near and far clipping distances for ray sampling based on camera matrix and volume bounds in IJK space.

Parameters:

camera_matrixtorch.Tensor (B, 4, 4) or (4, 4)

Batched or single 4x4 camera view matrix in RAS space Can be world-to-camera or camera-to-world (will be auto-detected)

volume_shapetuple of 3 ints

Volume dimensions in IJK space (I, J, K)

ijk2rastorch.Tensor (4, 4)

Transformation matrix from IJK to RAS coordinates

marginfloat, optional

Safety margin to add to near/far distances in mm (default: 30.0). Positive values expand the clipping range.

Returns:

neartorch.Tensor (B,) or float

Near clipping distance in mm (distance from camera to closest volume point)

fartorch.Tensor (B,) or float

Far clipping distance in mm (distance from camera to farthest volume point)

forward(vol, view_mat, ras2ijk, tile_h=256, tile_w=256, triton=False)[source]

Render one or more views of a (batched) 3D attenuation volume.

Parameters:

  • vol (torch.Tensor) – Attenuation volume of shape (B, C, D, H, W) in IJK voxel space. The affine ras2ijk is what makes the volume RAS-aligned at runtime.

  • view_mat (torch.Tensor) – Camera-to-world view matrices in RAS, shape (B, 4, 4) or (N, 4, 4) where N = B * views_per_vol. As a special case, N == 1 is broadcast to B. The clipping-distance computation also accepts world-to-camera matrices and auto-inverts them.

  • ras2ijk (torch.Tensor) – 4x4 affine mapping RAS to IJK coordinates.

  • tile_h (int) – Tile height used by the Python-loop backend to bound peak memory. Ignored if the image is smaller than (2*tile_h, 2*tile_w).

  • tile_w (int) – Tile width used by the Python-loop backend. See tile_h.

  • triton (bool) – If True, use the fused CUDA/Triton renderer (flashdrr.rendering.FusedVolumeRenderer); otherwise use the tiled Python-loop renderer.

Returns:

Rendered projections of shape (N, C, H, W).

Return type:

torch.Tensor

set_fov(fov)[source]

Configure the camera’s vertical field of view and refresh the per-pixel ray direction buffer self.dirs_cam.

The buffer is built on CPU in the image plane, normalized, then moved to the module’s current device.

Parameters:

fov (float) – Vertical field of view in degrees.

Return type:

None

FusedVolumeRenderer

class flashdrr.rendering.FusedVolumeRenderer[source]

Bases: Module

__init__(ray_samples, density_factor=1.0, apply_poisson_fn=None)[source]
Parameters:

  • ray_samples (int)

  • density_factor (float)

forward(density, view_mat, near, far, ras2ijk, vol_shape, dirs_cam)
Parameters:
Return type:

torch.Tensor

Scatter & context blocks

These building blocks are used when a physically motivated scatter term should be added to the primary Beer–Lambert image.

DepthAwareScatter

class flashdrr.rendering.DepthAwareScatter[source]

Bases: Module

Prototype depth-aware scatter estimator for DRR rendering.

Warning

Prototype, not fully functional. This module is an experimental design and is not yet validated end-to-end. The scattering model, gating, and integration with the main raycaster are subject to change. Do not rely on its outputs for production use.

Given per-sample attenuation mu along each ray, it computes a Beer-Lambert primary transmission, accumulates a depth-dependent scatter source, mixes it with an ASPP-conditioned gate, and adds a bounded scatter contribution to the primary image.

Parameters:

  • in_ch – Number of scatter channels (C) in the input volume.

  • base_ch – Hidden channel width of the conditioning/ASPP branch.

  • alpha_max – Upper bound on the learned scatter gating weight applied to the scatter map before it is added to the primary image.

__init__(in_ch, base_ch=32, alpha_max=0.4)[source]
forward(mu, dz=1.0)[source]

Estimate primary and scatter contributions along rays.

Parameters:

  • mu – Per-sample attenuation, shape (B, C, D, H, W) where D is the number of samples along each ray.

  • dz – Scalar sample spacing along the ray (same units as volume voxel sizes, e.g. mm). Defaults to 1.0.

Returns:

  • I_out ((B, C, H, W)): Primary transmission plus the bounded, gated scatter map.

  • I_primary ((B, C, H, W)): Beer-Lambert primary transmission at the exit of the ray (pre-scatter).

  • scatter_map ((B, 1, H, W)): Learned per-pixel scatter contribution before the alpha gate.

  • alpha ((B, 1, H, W)): Per-pixel gating weight in [0, alpha_max] applied to scatter_map.

Return type:

Tuple of

ASPP

class flashdrr.rendering.ASPP[source]

Bases: Module

Atrous Spatial Pyramid Pooling (ASPP) module.

Applies parallel 2D atrous (dilated) convolutions with multiple dilation rates, concatenates their ReLU-activated outputs, and projects the concatenated feature map back to out_ch channels with a 1x1 conv.

Parameters:

  • in_ch – Number of input channels.

  • out_ch – Number of output channels for each branch and the final projection.

  • rates – Tuple of dilation rates for the parallel Conv2d branches. Each branch uses a 3x3 kernel with padding equal to its dilation.

__init__(in_ch, out_ch, rates=(1, 2, 4, 8))[source]
forward(x)[source]

Run the ASPP branches and project to out_ch channels.

Parameters:

x – Input feature map of shape (B, in_ch, H, W).

Returns:

Output feature map of shape (B, out_ch, H, W).

Return type:

torch.Tensor

Camera & geometry helpers

Helper functions to build view/projection matrices and sample realistic C-arm (X-ray gantry) poses. The helpers re-exported from flashdrr.rendering are listed below; a few additional internal helpers (e.g. lookAt, homogenize_mat) live in flashdrr.rendering.raycast but are not part of the public API.

get_vtk_view_mat

Build a VTK-convention camera-to-world view matrix from a camera position, focal point and view-up vector in RAS coordinates.

flashdrr.rendering.get_vtk_view_mat(cam_pos, cam_focal, cam_viewup, device='cpu')

Build a VTK-compatible camera-to-world view matrix (4x4) from camera position, focal point, and view-up direction in RAS coordinates.

The resulting matrix maps camera-space points to world (RAS) space. The top-left 3x3 block has the camera axes as its columns, ordered [right, up, forward]; the fourth column is the camera position cam_pos; the bottom row is [0, 0, 0, 1]. forward points from cam_pos toward cam_focal.

Parameters:

  • cam_pos (Tuple[float]) – Camera/source position in RAS (mm), shape (3,).

  • cam_focal (Tuple[float]) – Camera focal (look-at) point in RAS (mm), shape (3,).

  • cam_viewup (Tuple[float]) – World-space up direction in RAS, shape (3,). Need not be unit length.

  • device (str) – Device on which to allocate the returned tensor. Defaults to 'cpu'.

Returns:

(4, 4) camera-to-world view matrix as described above.

Return type:

torch.Tensor

get_random_carm_views

Sample random C-arm geometry (SID, AP angle, lateral angle, table translation) and return a batch of VTK view matrices.

flashdrr.rendering.get_random_carm_views(n_views, sid_range, ap_range, lat_range, si_range, center)

Randomly sample C-arm view matrices by uniform sampling of the underlying geometric parameters.

For each of n_views samples, source-to-image distance (SID), AP angle, lateral angle, and table SI translation are drawn uniformly from the corresponding ranges, converted to a camera (position, focal, up) tuple via carm_to_camera_params(), and assembled into a 4x4 camera-to- world view matrix via get_vtk_view_mat(). The matrices are stacked along a new leading dimension.

Parameters:

n_viewsint

Number of view matrices to sample.

sid_rangetuple of 2 floats

(min_sid, max_sid) in mm. Typical range: (900, 1200)

ap_rangetuple of 2 floats

(min_ap, max_ap) in degrees. Typical range: (-40, 40) for cranial/caudal

lat_rangetuple of 2 floats

(min_lat, max_lat) in degrees. Typical range: (-90, 90) for LAO/RAO

si_rangetuple of 2 floats

(min_si, max_si) in mm. Typical range: (-100, 100) for table translation

centertuple or list of 3 floats

Volume center in RAS coordinates (mm), passed to carm_to_camera_params().

Returns:

viewstorch.Tensor (n_views, 4, 4)

Stacked camera-to-world view matrices in RAS space, on CPU.

Notes:

Distribution choice rationale:

  • SID: uniform is appropriate — mechanical constraint, no preferred distance.

  • AP angle: uniform is reasonable for training, though clinical use shows bias toward AP (0°), cranial (15-30°), and steep cranial (>30°).

  • Lateral angle: uniform works, but clinical practice favors AP (0°), RAO 30°, and LAO 30-45° views.

  • Table SI: uniform is appropriate — depends on anatomy region of interest.

For more realistic clinical distributions, consider adding bias toward common views (0°, ±30°, ±45°) or using a mixture of uniforms/peaked distributions.

carm_to_camera_params

Convert C-arm gantry parameters into a camera position, look-at point and view-up vector in RAS coordinates.

flashdrr.rendering.carm_to_camera_params(sid, ap_angle, lat_angle, center_ras, table_si=0.0)

Convert C-arm position parameters to camera parameters.

Parameters:

sidfloat

Source-to-Image Distance in mm (distance from X-ray source to detector)

ap_anglefloat

Anteroposterior (AP) angle in degrees - 0° = AP view (source in front, looking posterior) - Positive = caudal angulation (source tilts inferiorly) - Negative = cranial angulation (source tilts superiorly)

lat_anglefloat

Lateral angle in degrees - 0° = AP view - Positive = RAO (source moves to patient’s right) - Negative = LAO (source moves to patient’s left)

center_rastuple or list of 3 floats

Center point of the CT scan in RAS coordinates (mm) RAS = Right, Anterior, Superior

table_sifloat, optional

Table translation in superior-inferior direction in mm (default: 0.0) - Positive = table moves superior (head up) - Negative = table moves inferior (head down)

Returns:

cam_posnumpy array (3,)

Camera/source position in RAS coordinates

look_atnumpy array (3,)

Look-at point (focal point) in RAS coordinates

look_upnumpy array (3,)

Up vector for camera orientation

Transfer functions

piecewise_linear_channelwise

Per-channel piecewise-linear transfer function, typically used to map CT values (normalized Hounsfield units) to attenuation/density or RGBA before raycasting.

flashdrr.rendering.piecewise_linear_channelwise(x, xp, yp)[source]

Apply a per-channel piecewise linear function to input x. Handles extrapolation beyond xp bounds by extending the first/last segment.

Parameters:

  • x (Tensor) – Input (B, C, H, W) or (B, C, D, H, W)

  • xp (Tensor) – (C, K) sorted x keypoints per channel

  • yp (Tensor) – (C, K) y values per channel at keypoints

Returns:

Output of same shape as x

Return type:

Tensor