Differentiable Rasterizer

A high-performance, differentiable rasterizer that maps 2D and 3D point clouds to 2D grids using bilinear interpolation. Supports both a PyTorch-only implementation and an optional custom CUDA backend for better memory usage.

---

Installation

This package depends on PyTorch. Please install the version of PyTorch that matches your system and CUDA version before installing differentiable-rasterizer. Check the official PyTorch installation instructions for details.

Then clone this repository and install:

git clone https://github.com/microscopic-image-analysis/differentiable-rasterizer
cd differentiable-rasterizer
pip install -e .

Test it with:

python test_rasterizer.py

You should see output similar to:

Failed to import raster_cuda_impl. Falling back to torch version.

Testing rasterizer on grid (64, 64)
GPU used! outputs match CPU outputs!

Testing rasterizer on grid (128, 64)
GPU used! outputs match CPU outputs!

Testing rasterizer on grid (64, 1)
GPU used! outputs match CPU outputs!

All tests passed!

As you notice, the CUDA version of the raster function is not yet compiled. It is now falling back to the torch version, which should be fine for basic use cases. The CUDA version implements a custom CUDA kernel and uses orders of magnitude less memory.

To compile it, you need the CUDA Toolkit installed (not just the CUDA runtime that PyTorch uses), which provides the nvcc compiler. The toolkit is large and version-specific. You can download it here. Then run:

python build_cuda.py

and test with:

python test_rasterizer.py

You should now see output similar to:

CUDA extension detected!

Testing rasterizer on grid (64, 64)
GPU used! outputs match CPU outputs!

Testing rasterizer on grid (128, 64)
GPU used! outputs match CPU outputs!

Testing rasterizer on grid (64, 1)
GPU used! outputs match CPU outputs!

All tests passed!

Note: the build script places the final binary (.so for Linux) directly in where/you/cloned/the/repo/differentiable-rasterizer/src/rasterizer. Therefore, if you did not use an editable install (-e flag) earlier, you will need to move the binary to path/to/your/python/site-pakages/differentiable-rasterizer/src/rasterizer.

---

Usage

Low-level raster function (2d to 2d)

Use the raster function to convert 2D points to a 2D raster grid. Supports both CPU and GPU.

import torch
from rasterizer import raster

# Example inputs
points = torch.rand(2, 1000, 2) * 2 - 1    # shape (B, N, 2) , values in [-1, 1]
weights = torch.rand(1000)                 # shape (N,)
grid_size = (64, 64)

# Rasterize on CPU
images_cpu = raster(points, weights, grid_size, device="cpu")

# Rasterize on GPU
images_gpu = raster(points.to("cuda"), weights.to("cuda"), grid_size, device="cuda")

High-level class (3d to 2d)

The DifferentiableRasterizer can also be used as a high-level interface to rasterize 3D points onto a 2D grid. Points are first rotated in 3d space, then projected along the z-axis, and finally mapped to a 2d grid. Optionally, Gaussain smoothing is applied when a bandwidth parameter is provided.

import torch
from scipy.spatial.transform import Rotation as R
from rasterizer import DifferentiableRasterizer

H = 256 # height of images
W = 256 # width of images
N = 1000 # number of points
B = 10  # batch dimension

# input data
points = torch.rand(N, 3, dtype=torch.float32, device="cuda") * 2 - 1
weights = torch.ones(N, dtype=torch.float32, device="cuda")
rotations = torch.tensor(R.random(B).as_matrix(), dtype=torch.float32, device="cuda")  # shape: (B, 3, 3)
shifts = torch.rand(B, 2, dtype=torch.float32, device="cuda") * 0.1

# Initialize rasterizer
rasterizer = DifferentiableRasterizer(
    grid_size=(256, 256),
    kernel_size=5,     # kernel_size used for Gaussian smoothing            
    device="cuda"
)


# Rasterize 3D points to 2D grids
projections = rasterizer(
    points=points,
    weights=weights,
    rotations=rotations,
    translations=shifts,
    bandwidth=1.5           # in units of pixels. use None for no smoothing.
)

# `projections` has shape (B, H, W)