File size: 4,716 Bytes
e3043d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """
Pure-PyTorch differentiable point-based Gaussian renderer.
No CUDA compilation needed — works with standard PyTorch CUDA ops.
Simplified splat renderer:
- Projects 3D points to 2D screen space
- Renders each point as a 2D Gaussian splat
- Supports per-point color, opacity, and extra feature channels
- Differentiable w.r.t. point positions, colors, opacities, scales
"""
import torch
import torch.nn.functional as F
def project_points(means3d, viewmat, K, height, width):
"""Project 3D points to 2D screen coordinates."""
R = viewmat[:3, :3]
t = viewmat[:3, 3]
cam_points = means3d @ R.T + t
depths = cam_points[:, 2]
visible = depths > 0.01
fx, fy = K[0, 0], K[1, 1]
cx, cy = K[0, 2], K[1, 2]
safe_depths = torch.where(depths > 0.01, depths, torch.ones_like(depths))
points_2d = torch.stack([
fx * cam_points[:, 0] / safe_depths + cx,
fy * cam_points[:, 1] / safe_depths + cy,
], dim=-1)
return points_2d, depths, visible
def render_splat(points_2d, depths, colors, opacities, scales, height, width,
radius_factor=3.0):
"""Render points as 2D Gaussian splats using sorted alpha blending.
Fully functional (no in-place ops) for differentiability.
Gradients flow through colors, opacities, scales, and point positions.
The alpha compositing buffer is detached (controls order only).
"""
N = points_2d.shape[0]
C = colors.shape[1]
device = points_2d.device
margin = 50
valid = (depths > 0.01) & \
(points_2d[:, 0] > -margin) & (points_2d[:, 0] < width + margin) & \
(points_2d[:, 1] > -margin) & (points_2d[:, 1] < height + margin)
if valid.sum() == 0:
return torch.zeros(height, width, C, device=device)
pts = points_2d[valid]
cols = colors[valid]
ops = opacities[valid]
scs = scales[valid]
deps = depths[valid]
# Sort by depth (back-to-front)
sort_idx = torch.argsort(deps, descending=True)
pts = pts[sort_idx]
cols = cols[sort_idx]
ops = ops[sort_idx]
scs = scs[sort_idx]
radii = torch.clamp(scs * radius_factor, min=0.5, max=20.0)
# Detached alpha compositing buffer — purely for ordering
alpha_accum = torch.ones(height, width, device=device)
# Collect all contributions, then sum functionally at the end
image = torch.zeros(height, width, C, device=device)
M = pts.shape[0]
for i in range(M):
px = pts[i, 0]
py = pts[i, 1]
r_int = int(radii[i].item()) + 1
if r_int < 1:
continue
x0 = max(0, int(px.item()) - r_int)
x1 = min(width, int(px.item()) + r_int + 1)
y0 = max(0, int(py.item()) - r_int)
y1 = min(height, int(py.item()) + r_int + 1)
if x1 <= x0 or y1 <= y0:
continue
ys = torch.arange(y0, y1, device=device).float()
xs = torch.arange(x0, x1, device=device).float()
yy, xx = torch.meshgrid(ys, xs, indexing='ij')
sigma = radii[i]
gauss = torch.exp(-((xx - px) ** 2 + (yy - py) ** 2) / (2 * sigma ** 2 + 1e-8))
alpha = gauss * ops[i] # differentiable w.r.t. ops, sigma, px, py
# Read current alpha (detached)
current_alpha = alpha_accum[y0:y1, x0:x1].detach()
blend_alpha = alpha * current_alpha # differentiable
# Add contribution to image via functional padding (no in-place)
contrib = blend_alpha.unsqueeze(-1) * cols[i].unsqueeze(0).unsqueeze(0)
# contrib: [h_region, w_region, C]
pad_top = y0
pad_bottom = height - y1
pad_left = x0
pad_right = width - x1
padded = F.pad(contrib.permute(2, 0, 1),
(pad_left, pad_right, pad_top, pad_bottom),
mode='constant', value=0).permute(1, 2, 0)
image = image + padded
# Update detached alpha buffer (clone to avoid in-place modification)
new_region = current_alpha * (1 - alpha.detach())
alpha_new = alpha_accum.clone()
alpha_new[y0:y1, x0:x1] = new_region
alpha_accum = alpha_new.detach()
return image
def rasterize_gaussians(means3d, quats, scales, opacities, colors,
viewmat, K, height, width):
"""Main rasterization entry point."""
points_2d, depths, visible = project_points(means3d, viewmat, K, height, width)
fx = K[0, 0]
safe_depths = torch.where(depths > 0.01, depths, torch.ones_like(depths))
screen_scales = (fx * scales.mean(dim=-1) / safe_depths).clamp(min=0.5, max=50.0)
image = render_splat(points_2d, depths, colors, opacities, screen_scales,
height, width)
return image
|