Spaces:
Running on Zero
Running on Zero
| from math import isqrt | |
| from typing import Literal | |
| import torch | |
| from einops import rearrange, repeat | |
| from jaxtyping import Float | |
| from torch import Tensor | |
| from gsplat.rendering import rasterization | |
| GsplatRenderMode = Literal["RGB", "RGB+D", "RGB+ED"] | |
| GsplatRasterizeMode = Literal["classic", "antialiased"] | |
| def normalized_intrinsics_to_pixel_K( | |
| intrinsics: Float[Tensor, "batch view 3 3"], | |
| image_shape: tuple[int, int], | |
| ) -> Float[Tensor, "batch view 3 3"]: | |
| h, w = image_shape | |
| K = intrinsics.clone().float() | |
| K[..., 0, 0] *= w | |
| K[..., 1, 1] *= h | |
| K[..., 0, 2] *= w | |
| K[..., 1, 2] *= h | |
| return K | |
| def _identity_quats_like( | |
| means: Float[Tensor, "batch gaussian 3"], | |
| ) -> Float[Tensor, "batch gaussian 4"]: | |
| quats = torch.zeros((*means.shape[:-1], 4), device=means.device, dtype=means.dtype) | |
| # gsplat uses wxyz quaternions. | |
| quats[..., 0] = 1.0 | |
| return quats | |
| def _xyzw_to_wxyz( | |
| quats: Float[Tensor, "batch gaussian 4"], | |
| ) -> Float[Tensor, "batch gaussian 4"]: | |
| return torch.cat((quats[..., 3:4], quats[..., :3]), dim=-1) | |
| def render_gsplat( | |
| extrinsics: Float[Tensor, "batch view 4 4"], | |
| intrinsics: Float[Tensor, "batch view 3 3"], | |
| near: Float[Tensor, "batch view"], | |
| far: Float[Tensor, "batch view"], | |
| image_shape: tuple[int, int], | |
| background_color: Float[Tensor, "batch 3"], | |
| gaussian_means: Float[Tensor, "batch gaussian 3"], | |
| gaussian_covariances: Float[Tensor, "batch gaussian 3 3"] | None, | |
| gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], | |
| gaussian_opacities: Float[Tensor, "batch gaussian"], | |
| gaussian_scales: Float[Tensor, "batch gaussian 3"] | None = None, | |
| gaussian_rotations: Float[Tensor, "batch gaussian 4"] | None = None, | |
| render_mode: GsplatRenderMode = "RGB", | |
| use_covariances: bool = True, | |
| eps2d: float = 0.1, | |
| packed: bool = True, | |
| sparse_grad: bool = False, | |
| absgrad: bool = False, | |
| rasterize_mode: GsplatRasterizeMode = "antialiased", | |
| channel_chunk: int = 32, | |
| scale_invariant: bool = False, | |
| ) -> tuple[ | |
| Float[Tensor, "batch view 3 height width"] | None, | |
| Float[Tensor, "batch view height width"] | None, | |
| Float[Tensor, "batch view height width"], | |
| ]: | |
| b, v, _, _ = extrinsics.shape | |
| h, w = image_shape | |
| extrinsics = extrinsics.float() | |
| intrinsics = intrinsics.float() | |
| near = near.float() | |
| far = far.float() | |
| background_color = background_color.float() | |
| gaussian_means = gaussian_means.float() | |
| gaussian_sh_coefficients = gaussian_sh_coefficients.float() | |
| gaussian_opacities = gaussian_opacities.float() | |
| covars = gaussian_covariances.float() if gaussian_covariances is not None else None | |
| scales = gaussian_scales.float() if gaussian_scales is not None else None | |
| # The adapter stores quaternions as xyzw, while gsplat expects wxyz. | |
| quats = ( | |
| _xyzw_to_wxyz(gaussian_rotations.float()) | |
| if gaussian_rotations is not None | |
| else None | |
| ) | |
| if use_covariances: | |
| if covars is None: | |
| raise ValueError("gsplat covariance mode requires gaussians.covariances.") | |
| if scales is None: | |
| scales = torch.ones_like(gaussian_means) | |
| if quats is None: | |
| quats = _identity_quats_like(gaussian_means) | |
| else: | |
| covars = None | |
| if scales is None or quats is None: | |
| raise ValueError( | |
| "gsplat scale/quaternion mode requires gaussians.scales and " | |
| "gaussians.rotations. Prefer use_covariances=true." | |
| ) | |
| if scale_invariant: | |
| scale = 1.0 / near[:, :1].clamp_min(1e-8) | |
| extrinsics = extrinsics.clone() | |
| extrinsics[..., :3, 3] = extrinsics[..., :3, 3] * scale[..., None] | |
| gaussian_means = gaussian_means * scale[:, :, None] | |
| near = near * scale | |
| far = far * scale | |
| if covars is not None: | |
| covars = covars * (scale[:, :, None, None] ** 2) | |
| elif scales is not None: | |
| scales = scales * scale[:, :, None] | |
| Ks = normalized_intrinsics_to_pixel_K(intrinsics, image_shape) | |
| viewmats = torch.linalg.inv(extrinsics) | |
| sh_degree = isqrt(gaussian_sh_coefficients.shape[-1]) - 1 | |
| shs = rearrange( | |
| gaussian_sh_coefficients, "b g c n -> b g n c" | |
| ).contiguous() | |
| render_rgb = "RGB" in render_mode | |
| render_depth = "D" in render_mode | |
| # gsplat's packed path has limited background broadcasting across render | |
| # modes. Our datasets use black backgrounds, so omitting the background keeps | |
| # the memory-efficient packed path available without changing semantics. | |
| backgrounds = ( | |
| None | |
| if packed | |
| else repeat(background_color, "b c -> b v c", v=v).to(Ks) | |
| ) | |
| render_out, alpha, _ = rasterization( | |
| means=gaussian_means, | |
| quats=quats, | |
| scales=scales, | |
| opacities=gaussian_opacities, | |
| colors=shs, | |
| viewmats=viewmats, | |
| Ks=Ks, | |
| width=w, | |
| height=h, | |
| near_plane=float(near.detach().min().item()), | |
| far_plane=float(far.detach().max().item()), | |
| eps2d=eps2d, | |
| sh_degree=sh_degree, | |
| packed=packed, | |
| backgrounds=backgrounds, | |
| render_mode=render_mode, | |
| sparse_grad=sparse_grad, | |
| absgrad=absgrad, | |
| rasterize_mode=rasterize_mode, | |
| channel_chunk=channel_chunk, | |
| covars=covars, | |
| ) | |
| color = ( | |
| rearrange(render_out[..., :3], "b v h w c -> b v c h w") | |
| if render_rgb | |
| else None | |
| ) | |
| depth = render_out[..., -1] if render_depth else None | |
| return color, depth, alpha.squeeze(-1) | |
| def render_depth_gsplat( | |
| extrinsics: Float[Tensor, "batch view 4 4"], | |
| intrinsics: Float[Tensor, "batch view 3 3"], | |
| near: Float[Tensor, "batch view"], | |
| far: Float[Tensor, "batch view"], | |
| image_shape: tuple[int, int], | |
| background_color: Float[Tensor, "batch 3"], | |
| gaussian_means: Float[Tensor, "batch gaussian 3"], | |
| gaussian_covariances: Float[Tensor, "batch gaussian 3 3"] | None, | |
| gaussian_sh_coefficients: Float[Tensor, "batch gaussian 3 d_sh"], | |
| gaussian_opacities: Float[Tensor, "batch gaussian"], | |
| gaussian_scales: Float[Tensor, "batch gaussian 3"] | None = None, | |
| gaussian_rotations: Float[Tensor, "batch gaussian 4"] | None = None, | |
| mode: Literal["depth", "disparity", "relative_disparity", "log"] = "depth", | |
| use_covariances: bool = True, | |
| eps2d: float = 0.1, | |
| packed: bool = True, | |
| sparse_grad: bool = False, | |
| absgrad: bool = False, | |
| rasterize_mode: GsplatRasterizeMode = "antialiased", | |
| channel_chunk: int = 32, | |
| scale_invariant: bool = False, | |
| ) -> Float[Tensor, "batch view height width"]: | |
| render_mode = "RGB+ED" if mode == "relative_disparity" else "RGB+D" | |
| _, depth, _ = render_gsplat( | |
| extrinsics, | |
| intrinsics, | |
| near, | |
| far, | |
| image_shape, | |
| background_color, | |
| gaussian_means, | |
| gaussian_covariances, | |
| gaussian_sh_coefficients, | |
| gaussian_opacities, | |
| gaussian_scales=gaussian_scales, | |
| gaussian_rotations=gaussian_rotations, | |
| render_mode=render_mode, | |
| use_covariances=use_covariances, | |
| eps2d=eps2d, | |
| packed=packed, | |
| sparse_grad=sparse_grad, | |
| absgrad=absgrad, | |
| rasterize_mode=rasterize_mode, | |
| channel_chunk=channel_chunk, | |
| scale_invariant=scale_invariant, | |
| ) | |
| assert depth is not None | |
| if mode in ("disparity", "relative_disparity"): | |
| depth = 1.0 / depth.clamp_min(1e-8) | |
| elif mode == "log": | |
| depth = depth.clamp_min(near[..., None, None]).clamp_max( | |
| far[..., None, None] | |
| ).log() | |
| return depth | |