Spaces:
Running on Zero
Running on Zero
| from dataclasses import dataclass | |
| from typing import Literal | |
| import torch | |
| from jaxtyping import Float | |
| from torch import Tensor | |
| from ...dataset import DatasetCfg | |
| from ..types import Gaussians | |
| from .decoder import Decoder, DecoderOutput | |
| from .gsplat_splatting import ( | |
| GsplatRasterizeMode, | |
| GsplatRenderMode, | |
| render_depth_gsplat, | |
| render_gsplat, | |
| ) | |
| class DecoderSplattingGsplatCfg: | |
| name: Literal["splatting_gsplat", "gsplat"] | |
| render_mode: GsplatRenderMode = "RGB" | |
| rasterize_mode: GsplatRasterizeMode = "antialiased" | |
| use_covariances: bool = True | |
| packed: bool = True | |
| eps2d: float = 0.1 | |
| sparse_grad: bool = False | |
| absgrad: bool = False | |
| channel_chunk: int = 32 | |
| scale_invariant: bool = False | |
| class DecoderSplattingGsplat(Decoder[DecoderSplattingGsplatCfg]): | |
| background_color: Float[Tensor, "3"] | |
| def __init__( | |
| self, | |
| cfg: DecoderSplattingGsplatCfg, | |
| dataset_cfg: DatasetCfg, | |
| ) -> None: | |
| super().__init__(cfg, dataset_cfg) | |
| self.register_buffer( | |
| "background_color", | |
| torch.tensor(dataset_cfg.background_color, dtype=torch.float32), | |
| persistent=False, | |
| ) | |
| def forward( | |
| self, | |
| gaussians: Gaussians, | |
| 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], | |
| depth_mode=None, | |
| ) -> DecoderOutput: | |
| render_mode = self.cfg.render_mode | |
| if depth_mode is not None: | |
| render_mode = "RGB+ED" if depth_mode == "relative_disparity" else "RGB+D" | |
| color, depth, accumulated_alpha = render_gsplat( | |
| extrinsics, | |
| intrinsics, | |
| near, | |
| far, | |
| image_shape, | |
| self.background_color[None].expand(extrinsics.shape[0], -1), | |
| gaussians.means, | |
| gaussians.covariances, | |
| gaussians.harmonics, | |
| gaussians.opacities, | |
| gaussian_scales=getattr(gaussians, "scales", None), | |
| gaussian_rotations=getattr( | |
| gaussians, | |
| "rotations_unnorm", | |
| getattr(gaussians, "rotations", None), | |
| ), | |
| render_mode=render_mode, | |
| use_covariances=self.cfg.use_covariances, | |
| eps2d=self.cfg.eps2d, | |
| packed=self.cfg.packed, | |
| sparse_grad=self.cfg.sparse_grad, | |
| absgrad=self.cfg.absgrad, | |
| rasterize_mode=self.cfg.rasterize_mode, | |
| channel_chunk=self.cfg.channel_chunk, | |
| scale_invariant=self.cfg.scale_invariant, | |
| ) | |
| assert color is not None | |
| if depth_mode in ("disparity", "relative_disparity"): | |
| assert depth is not None | |
| depth = 1.0 / depth.clamp_min(1e-8) | |
| elif depth_mode == "log": | |
| assert depth is not None | |
| depth = depth.clamp_min(near[..., None, None]).clamp_max( | |
| far[..., None, None] | |
| ).log() | |
| elif depth_mode is None and "D" not in render_mode: | |
| depth = None | |
| return DecoderOutput( | |
| color, | |
| depth, | |
| accumulated_alpha=accumulated_alpha, | |
| ) | |
| def render_depth( | |
| self, | |
| gaussians: Gaussians, | |
| 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], | |
| mode="depth", | |
| ) -> Float[Tensor, "batch view height width"]: | |
| return render_depth_gsplat( | |
| extrinsics, | |
| intrinsics, | |
| near, | |
| far, | |
| image_shape, | |
| self.background_color[None].expand(extrinsics.shape[0], -1), | |
| gaussians.means, | |
| gaussians.covariances, | |
| gaussians.harmonics, | |
| gaussians.opacities, | |
| gaussian_scales=getattr(gaussians, "scales", None), | |
| gaussian_rotations=getattr( | |
| gaussians, | |
| "rotations_unnorm", | |
| getattr(gaussians, "rotations", None), | |
| ), | |
| mode=mode, | |
| use_covariances=self.cfg.use_covariances, | |
| eps2d=self.cfg.eps2d, | |
| packed=self.cfg.packed, | |
| sparse_grad=self.cfg.sparse_grad, | |
| absgrad=self.cfg.absgrad, | |
| rasterize_mode=self.cfg.rasterize_mode, | |
| channel_chunk=self.cfg.channel_chunk, | |
| scale_invariant=self.cfg.scale_invariant, | |
| ) | |