cad-3d-reconstruction / code /simple_renderer.py

Upload code/simple_renderer.py with huggingface_hub

e3043d0 verified 2 months ago

4.72 kB

"""
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

	"""
	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