diff --git a/gsplat/rendering.py b/gsplat/rendering.py index 73ac13b0f..0e4456abd 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -2193,7 +2193,9 @@ def rasterization_2dgs( quats: Tensor, # [..., N, 4] scales: Tensor, # [..., N, 3] opacities: Tensor, # [..., N] - colors: Tensor, # [..., (C,) N, D] for post-activation colors, or [N, K, D] for SH coefficients + colors: Optional[ + Tensor + ], # [..., (C,) N, D] for post-activation colors, or [N, K, D] for SH coefficients viewmats: Tensor, # [..., C, 4, 4] Ks: Tensor, # [..., C, 3, 3] width: int, @@ -2211,6 +2213,9 @@ def rasterization_2dgs( absgrad: bool = False, distloss: bool = False, depth_mode: Literal["expected", "median"] = "expected", + extra_signals: Optional[Tensor] = None, + extra_signals_sh_degree: Optional[int] = None, + channel_chunk: int = 32, ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Dict]: """Rasterize a set of 2D Gaussians (N) to a batch of image planes (C). @@ -2260,7 +2265,9 @@ def rasterization_2dgs( will be done looply in chunks. distloss: If true, use distortion regularization to get better geometry detail. depth_mode: render depth mode. Choose from expected depth and median depth. - + extra_signals: Auxiliary per-Gaussian signals, shaped [..., N, E] or [..., C, N, E]. + extra_signals_sh_degree: If set, extra_signals is interpreted as SH coefficients [N, K, E]. + channel_chunk: Number of feature channels to rasterize per CUDA call. Returns: A tuple: @@ -2322,7 +2329,8 @@ def rasterization_2dgs( C = viewmats.shape[-3] I = B * C device = means.device - channels = colors.shape[-1] + has_color = render_mode_has_color(render_mode) + channels = colors.shape[-1] if has_color else 0 assert means.shape == batch_dims + (N, 3), means.shape assert quats.shape == batch_dims + (N, 4), quats.shape @@ -2335,7 +2343,13 @@ def rasterization_2dgs( render_mode ), f"distloss requires depth rendering, but render mode is {render_mode}" - if sh_degree is None: + if colors is None and has_color: + raise ValueError( + f"colors must be provided when render_mode='{render_mode}' includes RGB." + ) + if colors is None and sh_degree is not None: + raise ValueError("sh_degree must be None when colors is None.") + if has_color and sh_degree is None: # treat colors as post-activation values, should be in shape [..., N, D] or [..., C, N, D] assert ( colors.dim() == num_batch_dims + 2 @@ -2344,12 +2358,32 @@ def rasterization_2dgs( colors.dim() == num_batch_dims + 3 and colors.shape[:-1] == batch_dims + (C, N) ), colors.shape - else: + elif has_color: # treat colors as SH coefficients, must be in shape [N, K, D]. # Allowing for activating partial SH bands. assert colors.dim() == 3 and colors.shape[0] == N, colors.shape assert (sh_degree + 1) ** 2 <= colors.shape[-2], colors.shape + if extra_signals is not None: + extra_signal_channels = extra_signals.shape[-1] + if extra_signals_sh_degree is None: + assert ( + extra_signals.dim() == num_batch_dims + 2 + and extra_signals.shape[:-1] == batch_dims + (N,) + ) or ( + extra_signals.dim() == num_batch_dims + 3 + and extra_signals.shape[:-1] == batch_dims + (C, N) + ), extra_signals.shape + else: + assert ( + extra_signals.dim() == 3 and extra_signals.shape[0] == N + ), extra_signals.shape + assert (extra_signals_sh_degree + 1) ** 2 <= extra_signals.shape[ + -2 + ], extra_signals.shape + else: + extra_signal_channels = 0 + # Compute Ray-Splat intersection transformation. proj_results = fully_fused_projection_2dgs( means, @@ -2410,67 +2444,220 @@ def rasterization_2dgs( isect_offsets = isect_offset_encode(isect_ids, I, tile_width, tile_height) isect_offsets = isect_offsets.reshape(batch_dims + (C, tile_height, tile_width)) - if sh_degree is not None: # SH coefficients - if packed: - dirs = compute_directions( + valid_gaussians = (radii > 0).all(dim=-1) + + feature_list = [] + + if has_color: + if sh_degree is not None: # SH coefficients + if packed: + dirs = compute_directions( + batch_dims, + means, + viewmats, + batch_ids, + camera_ids, + gaussian_ids, + indptr, + ) # [nnz, 3] + shs = colors[gaussian_ids] + else: + camtoworlds = torch.inverse(viewmats) + dirs = means[..., None, :, :] - camtoworlds[..., None, :3, 3] + shs = colors + + colors = spherical_harmonics(sh_degree, dirs, shs, masks=valid_gaussians) + colors = torch.clamp_min(colors + 0.5, 0.0) + else: + colors = normalize_features_layout( + colors, batch_dims, - means, - viewmats, + C, + colors.shape[-1:], batch_ids, camera_ids, gaussian_ids, - indptr, - ) # [nnz, 3] - # Gather per-Gaussian coeffs to match the [nnz, K, 3] dirs. - shs = colors[gaussian_ids] + ) + + feature_list.append(colors) + + render_extra_signal_layout = None + extra_signal_source = None + if extra_signals is not None: + if extra_signals_sh_degree is not None: + render_extra_signal_layout = "sh" + extra_signal_source = "sh_evaluated" + if packed: + dirs = compute_directions( + batch_dims, + means, + viewmats, + batch_ids, + camera_ids, + gaussian_ids, + indptr, + ) # [nnz, 3] + shs = extra_signals[gaussian_ids] + else: + camtoworlds = torch.inverse(viewmats) + dirs = means[..., None, :, :] - camtoworlds[..., None, :3, 3] + shs = extra_signals + + extra_signals = spherical_harmonics( + extra_signals_sh_degree, + dirs, + shs, + masks=valid_gaussians, + ) + extra_signals = extra_signals + 0.5 else: - camtoworlds = torch.inverse(viewmats) - dirs = means[..., None, :, :] - camtoworlds[..., None, :3, 3] - shs = colors - colors = spherical_harmonics( - sh_degree, dirs, shs, masks=(radii > 0).all(dim=-1) - ) # [nnz, 3] or [..., C, N, 3] - # make it apple-to-apple with Inria's CUDA Backend. - colors = torch.clamp_min(colors + 0.5, 0.0) + render_extra_signal_layout = ( + "per_gaussian" + if extra_signals.dim() == num_batch_dims + 2 + else "per_camera" + ) + extra_signal_source = "post_activation" + extra_signals = normalize_features_layout( + extra_signals, + batch_dims, + C, + extra_signals.shape[-1:], + batch_ids, + camera_ids, + gaussian_ids, + ) + + feature_list.append(extra_signals) + + proj_features = ( + torch.cat(feature_list, dim=-1) + if len(feature_list) > 1 + else feature_list[0] + if feature_list + else None + ) + + if backgrounds is not None and extra_signals is not None: + backgrounds = torch.cat( + [ + backgrounds, + torch.zeros( + batch_dims + (C, extra_signal_channels), + device=backgrounds.device, + dtype=backgrounds.dtype, + ), + ], + dim=-1, + ) # Rasterize to pixels if render_mode_has_depth_channel(render_mode) and render_mode_has_color( render_mode ): - colors = torch.cat((colors, depths[..., None]), dim=-1) + proj_features = torch.cat((proj_features, depths[..., None]), dim=-1) if backgrounds is not None: backgrounds = torch.cat( (backgrounds, torch.zeros_like(backgrounds[..., :1])), dim=-1 ) elif render_mode_has_only_depth_channel(render_mode): - colors = depths[..., None] + depth_channel = depths[..., None] + if proj_features is not None: + proj_features = torch.cat((proj_features, depth_channel), dim=-1) + else: + proj_features = depth_channel else: # RGB pass - ( - render_colors, - render_alphas, - render_normals, - render_distort, - render_median, - ) = rasterize_to_pixels_2dgs( - means2d, - ray_transforms, - colors, - opacities, - normals, - densify, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - backgrounds=backgrounds, - packed=packed, - absgrad=absgrad, - distloss=distloss, - ) + if proj_features.shape[-1] > channel_chunk: + n_chunks = (proj_features.shape[-1] + channel_chunk - 1) // channel_chunk + render_colors = [] + render_alphas = [] + render_normals = [] + render_distorts = [] + render_medians = [] + for i in range(n_chunks): + features_chunk = proj_features[ + ..., i * channel_chunk : (i + 1) * channel_chunk + ] + backgrounds_chunk = ( + backgrounds[..., i * channel_chunk : (i + 1) * channel_chunk] + if backgrounds is not None + else None + ) + ( + render_colors_, + render_alphas_, + render_normals_, + render_distort_, + render_median_, + ) = rasterize_to_pixels_2dgs( + means2d, + ray_transforms, + features_chunk, + opacities, + normals, + densify, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + backgrounds=backgrounds_chunk, + packed=packed, + absgrad=absgrad, + distloss=distloss, + ) + render_colors.append(render_colors_) + render_alphas.append(render_alphas_) + render_normals.append(render_normals_) + render_distorts.append(render_distort_) + render_medians.append(render_median_) + + render_colors = torch.cat(render_colors, dim=-1) + render_alphas = render_alphas[0] + render_normals = render_normals[0] + render_distort = render_distorts[0] + render_median = render_medians[0] + else: + ( + render_colors, + render_alphas, + render_normals, + render_distort, + render_median, + ) = rasterize_to_pixels_2dgs( + means2d, + ray_transforms, + proj_features, + opacities, + normals, + densify, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + backgrounds=backgrounds, + packed=packed, + absgrad=absgrad, + distloss=distloss, + ) + + render_extra_signals = None + if extra_signals is not None: + render_extra_signals = render_colors[ + ..., channels : channels + extra_signal_channels + ] + + if render_mode_has_depth_channel(render_mode): + render_depth = render_colors[..., -1:] + render_colors = torch.cat( + [render_colors[..., :channels], render_depth], dim=-1 + ) + else: + render_colors = render_colors[..., :channels] + render_normals_from_depth = None if render_mode_has_expected_depth(render_mode): # normalize the accumulated depth to get the expected depth @@ -2515,6 +2702,14 @@ def rasterization_2dgs( "gradient_2dgs": densify, # This holds the gradient used for densification for 2dgs } + if render_extra_signals is not None: + meta["render_extra_signals"] = render_extra_signals + meta["extra_signal_layout"] = render_extra_signal_layout + meta["extra_signal_channels"] = extra_signal_channels + meta["extra_signals_sh_degree"] = extra_signals_sh_degree + meta["extra_signal_source"] = extra_signal_source + meta["extra_signal_compositing"] = "alpha" + render_normals = torch.einsum( "...ij,...hwj->...hwi", torch.linalg.inv_ex(viewmats).inverse[..., :3, :3], diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 4ab94eeb3..6b19f32a2 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -74,6 +74,117 @@ def test_data(): } +@pytest.fixture +def deterministic_2dgs_data(): + torch.manual_seed(7) + + C = 2 + N = 6 + W, H = 32, 24 + signal_channels = 20 + means = torch.tensor( + [ + [-0.30, -0.20, 2.00], + [-0.10, 0.15, 2.20], + [0.12, -0.05, 2.40], + [0.28, 0.18, 2.60], + [-0.22, 0.24, 2.80], + [0.05, -0.28, 3.00], + ], + dtype=torch.float32, + device=device, + ) + quats = torch.tensor([[1.0, 0.0, 0.0, 0.0]] * N, dtype=torch.float32, device=device) + scales = torch.tensor([[0.18, 0.14, 1.0]] * N, dtype=torch.float32, device=device) + opacities = torch.linspace(0.20, 0.70, N, device=device) + colors = torch.linspace(0.05, 0.95, C * N * 3, device=device).reshape(C, N, 3) + viewmats = torch.eye(4, device=device).expand(C, 4, 4).clone() + viewmats[1, 0, 3] = 0.05 + + Ks = ( + torch.tensor( + [ + [float(W), 0.0, W / 2.0], + [0.0, float(W), H / 2.0], + [0.0, 0.0, 1.0], + ], + device=device, + ) + .expand(C, 3, 3) + .clone() + ) + + extra_signals = torch.linspace( + 0.10, 0.90, C * N * signal_channels, device=device + ).reshape(C, N, signal_channels) + + return { + "means": means, + "quats": quats, + "scales": scales, + "opacities": opacities, + "colors": colors, + "viewmats": viewmats, + "Ks": Ks, + "extra_signals": extra_signals, + "width": W, + "height": H, + "n_cameras": C, + "n_gaussians": N, + "signal_channels": signal_channels, + } + + +def _assert_2dgs_outputs_close(actual, expected): + ( + actual_render_colors, + actual_render_alphas, + actual_normals, + actual_surf_normals, + ) = actual[:4] + actual_render_distort, actual_render_median, actual_meta = actual[4:] + + expected_render_colors, expected_render_alphas, expected_normals = expected[:3] + expected_surf_normals, expected_render_distort, expected_render_median = expected[ + 3:6 + ] + expected_meta = expected[6] + + torch.testing.assert_close(actual_render_colors, expected_render_colors) + torch.testing.assert_close(actual_render_alphas, expected_render_alphas) + torch.testing.assert_close(actual_normals, expected_normals) + torch.testing.assert_close(actual_render_distort, expected_render_distort) + torch.testing.assert_close(actual_render_median, expected_render_median) + + if actual_surf_normals is None: + assert expected_surf_normals is None + else: + torch.testing.assert_close(actual_surf_normals, expected_surf_normals) + + for key in ( + "radii", + "means2d", + "depths", + "ray_transforms", + "opacities", + "normals", + "isect_offsets", + "width", + "height", + "tile_size", + "n_cameras", + "render_distort", + "gradient_2dgs", + ): + assert key in actual_meta + assert key in expected_meta + + if isinstance(actual_meta[key], torch.Tensor): + assert actual_meta[key].shape == expected_meta[key].shape + else: + assert actual_meta[key] == expected_meta[key] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") @pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) @@ -485,6 +596,56 @@ def test_rasterize_to_pixels_2dgs( ) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize("packed", [False, True]) +def test_rasterization_2dgs_extra_signals_sh( + deterministic_2dgs_data, + packed: bool, +): + from gsplat.rendering import rasterization_2dgs + + sh_degree = 3 + signal_channels = 20 + K = (sh_degree + 1) ** 2 + extra_signals = torch.linspace( + 0.10, + 0.90, + deterministic_2dgs_data["n_gaussians"] * K * signal_channels, + device=device, + ).reshape(deterministic_2dgs_data["n_gaussians"], K, signal_channels) + + render_colors, _, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=deterministic_2dgs_data["colors"], + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode="RGB", + packed=packed, + extra_signals=extra_signals, + extra_signals_sh_degree=sh_degree, + ) + + assert render_colors.shape[-1] == 3 + assert meta["render_extra_signals"].shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + signal_channels, + ) + assert meta["extra_signal_layout"] == "sh" + assert meta["extra_signal_channels"] == signal_channels + assert meta["extra_signals_sh_degree"] == sh_degree + assert meta["extra_signal_source"] == "sh_evaluated" + assert meta["extra_signal_compositing"] == "alpha" + assert torch.isfinite(meta["render_extra_signals"]).all() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") @pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") @pytest.mark.parametrize("batch_dims", [(2,), (1, 2)]) @@ -521,6 +682,394 @@ def test_rasterization_packed_2dgs(test_data, batch_dims: Tuple[int, ...]): assert torch.isfinite(render_colors).all() +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize("render_mode", ["RGB", "RGB+D"]) +def test_rasterization_2dgs_extra_signals_none_matches_existing_path( + deterministic_2dgs_data, + render_mode: str, + packed: bool = False, +): + from gsplat.rendering import rasterization_2dgs + + inputs = { + "means": deterministic_2dgs_data["means"], + "quats": deterministic_2dgs_data["quats"], + "scales": deterministic_2dgs_data["scales"], + "opacities": deterministic_2dgs_data["opacities"], + "colors": deterministic_2dgs_data["colors"], + "viewmats": deterministic_2dgs_data["viewmats"], + "Ks": deterministic_2dgs_data["Ks"], + "width": deterministic_2dgs_data["width"], + "height": deterministic_2dgs_data["height"], + "render_mode": render_mode, + "packed": packed, + } + + expected = rasterization_2dgs(**inputs) + + actual = rasterization_2dgs( + **inputs, + extra_signals=None, + extra_signals_sh_degree=None, + ) + + _assert_2dgs_outputs_close(actual, expected) + + actual_meta = actual[-1] + assert "render_extra_signals" not in actual_meta + assert "extra_signal_layout" not in actual_meta + assert "extra_signal_channels" not in actual_meta + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize( + "render_mode,expected_render_channels", [("RGB", 3), ("RGB+D", 4)] +) +def test_rasterization_2dgs_extra_signals_post_activation_metadata_and_shape( + deterministic_2dgs_data, + render_mode: str, + expected_render_channels: int, +): + from gsplat.rendering import rasterization_2dgs + + render_colors, render_alphas, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=deterministic_2dgs_data["colors"], + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode=render_mode, + extra_signals=deterministic_2dgs_data["extra_signals"], + extra_signals_sh_degree=None, + ) + + assert render_colors.shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + expected_render_channels, + ) + assert render_alphas.shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + 1, + ) + + assert "render_extra_signals" in meta + assert meta["render_extra_signals"].shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + deterministic_2dgs_data["signal_channels"], + ) + + assert meta["extra_signal_layout"] == "per_camera" + assert meta["extra_signal_channels"] == deterministic_2dgs_data["signal_channels"] + assert meta["extra_signals_sh_degree"] is None + assert meta["extra_signal_source"] == "post_activation" + assert meta["extra_signal_compositing"] == "alpha" + + assert torch.isfinite(render_colors).all() + assert torch.isfinite(meta["render_extra_signals"]).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize( + "extra_signals,expected_layout", + [ + ("per_gaussian", "per_gaussian"), + ("per_camera", "per_camera"), + ], +) +def test_rasterization_2dgs_extra_signals_layout_metadata( + deterministic_2dgs_data, + extra_signals: str, + expected_layout: str, +): + from gsplat.rendering import rasterization_2dgs + + signal_input = ( + deterministic_2dgs_data["extra_signals"][0] + if extra_signals == "per_gaussian" + else deterministic_2dgs_data["extra_signals"] + ) + + _, _, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=deterministic_2dgs_data["colors"], + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode="RGB", + extra_signals=signal_input, + ) + + assert meta["extra_signal_layout"] == expected_layout + assert meta["extra_signal_channels"] == deterministic_2dgs_data["signal_channels"] + assert meta["render_extra_signals"].shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + deterministic_2dgs_data["signal_channels"], + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize("extra_signals", ["per_gaussian", "per_camera"]) +def test_rasterization_2dgs_extra_signals_packed_post_activation( + deterministic_2dgs_data, + extra_signals: str, +): + from gsplat.rendering import rasterization_2dgs + + signal_input = ( + deterministic_2dgs_data["extra_signals"][0] + if extra_signals == "per_gaussian" + else deterministic_2dgs_data["extra_signals"] + ) + + render_colors, _, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=deterministic_2dgs_data["colors"], + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode="RGB", + packed=True, + extra_signals=signal_input, + ) + + assert render_colors.shape[-1] == 3 + assert ( + meta["render_extra_signals"].shape[-1] + == deterministic_2dgs_data["signal_channels"] + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize("render_mode", ["D", "ED"]) +@pytest.mark.parametrize("packed", [False, True]) +def test_rasterization_2dgs_extra_signals_depth_only( + deterministic_2dgs_data, + render_mode: str, + packed: bool, +): + from gsplat.rendering import rasterization_2dgs + + render_colors, render_alphas, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=None, + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode=render_mode, + packed=packed, + extra_signals=deterministic_2dgs_data["extra_signals"], + ) + + assert render_colors.shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + 1, + ) + assert render_alphas.shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + 1, + ) + assert meta["render_extra_signals"].shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + deterministic_2dgs_data["signal_channels"], + ) + assert meta["extra_signal_channels"] == deterministic_2dgs_data["signal_channels"] + assert torch.isfinite(render_colors).all() + assert torch.isfinite(meta["render_extra_signals"]).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize("packed", [False, True]) +def test_rasterization_2dgs_extra_signals_channel_chunk( + deterministic_2dgs_data, + packed: bool, +): + from gsplat.rendering import rasterization_2dgs + + signal_channels = 4 + extra_signals = deterministic_2dgs_data["extra_signals"][..., :signal_channels] + + render_colors, _, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=deterministic_2dgs_data["colors"], + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode="RGB", + packed=packed, + extra_signals=extra_signals, + channel_chunk=4, + ) + + assert render_colors.shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + 3, + ) + assert meta["render_extra_signals"].shape == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + signal_channels, + ) + assert meta["extra_signal_channels"] == signal_channels + assert torch.isfinite(render_colors).all() + assert torch.isfinite(meta["render_extra_signals"]).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +@pytest.mark.parametrize("packed", [False, True]) +def test_rasterization_2dgs_extra_signals_gradients( + deterministic_2dgs_data, + packed: bool, +): + from gsplat.rendering import rasterization_2dgs + + means = deterministic_2dgs_data["means"].detach().clone().requires_grad_(True) + quats = deterministic_2dgs_data["quats"].detach().clone().requires_grad_(True) + scales = deterministic_2dgs_data["scales"].detach().clone().requires_grad_(True) + opacities = ( + deterministic_2dgs_data["opacities"].detach().clone().requires_grad_(True) + ) + colors = deterministic_2dgs_data["colors"].detach().clone().requires_grad_(True) + extra_signals = ( + deterministic_2dgs_data["extra_signals"].detach().clone().requires_grad_(True) + ) + + render_colors, render_alphas, _, _, _, _, meta = rasterization_2dgs( + means=means, + quats=quats, + scales=scales, + opacities=opacities, + colors=colors, + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode="RGB+D", + packed=packed, + extra_signals=extra_signals, + ) + + loss = ( + render_colors.sum() + render_alphas.sum() + meta["render_extra_signals"].sum() + ) + loss.backward() + + for name, tensor in [ + ("means", means), + ("quats", quats), + ("scales", scales), + ("opacities", opacities), + ("colors", colors), + ("extra_signals", extra_signals), + ]: + assert tensor.grad is not None, f"{name} should receive gradients" + assert torch.isfinite(tensor.grad).all(), f"{name} grad should be finite" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") +def test_rasterization_2dgs_extra_signals_deterministic_stats( + deterministic_2dgs_data, +): + from gsplat.rendering import rasterization_2dgs + + signal_channels = 4 + extra_signals = deterministic_2dgs_data["extra_signals"][..., :signal_channels] + + render_colors, render_alphas, _, _, _, _, meta = rasterization_2dgs( + means=deterministic_2dgs_data["means"], + quats=deterministic_2dgs_data["quats"], + scales=deterministic_2dgs_data["scales"], + opacities=deterministic_2dgs_data["opacities"], + colors=deterministic_2dgs_data["colors"], + viewmats=deterministic_2dgs_data["viewmats"], + Ks=deterministic_2dgs_data["Ks"], + width=deterministic_2dgs_data["width"], + height=deterministic_2dgs_data["height"], + render_mode="RGB+D", + extra_signals=extra_signals, + channel_chunk=4, + ) + + stats = { + "n_gaussians": deterministic_2dgs_data["n_gaussians"], + "signal_channels": signal_channels, + "render_shape": tuple(render_colors.shape), + "extra_shape": tuple(meta["render_extra_signals"].shape), + "alpha_sum": render_alphas.sum(), + "render_sum": render_colors.sum(), + "extra_sum": meta["render_extra_signals"].sum(), + } + + assert stats["n_gaussians"] == 6 + assert stats["signal_channels"] == 4 + assert stats["render_shape"] == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + 4, + ) + assert stats["extra_shape"] == ( + deterministic_2dgs_data["n_cameras"], + deterministic_2dgs_data["height"], + deterministic_2dgs_data["width"], + signal_channels, + ) + + torch.testing.assert_close( + stats["alpha_sum"], torch.tensor(114.67495727539062, device=device) + ) + torch.testing.assert_close( + stats["render_sum"], torch.tensor(462.40618896484375, device=device) + ) + torch.testing.assert_close( + stats["extra_sum"], torch.tensor(219.33282470703125, device=device) + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") @pytest.mark.skipif(not gsplat.has_2dgs(), reason="2DGS support wasn't built") def test_rasterization_packed_2dgs_pose_grad_large_nnz():