Spaces:
Running on Zero
Running on Zero
| from dataclasses import dataclass | |
| from typing import Literal | |
| import torch | |
| from einops import rearrange, repeat | |
| from jaxtyping import Float | |
| from torch import Tensor | |
| from ...dataset import DatasetCfg | |
| from ..types import Gaussians | |
| from .cuda_splatting import DepthRenderingMode, render_cuda, render_depth_cuda | |
| from .decoder import Decoder, DecoderOutput | |
| class DecoderSplattingCUDACfg: | |
| name: Literal["splatting_cuda"] | |
| class DecoderSplattingCUDA(Decoder[DecoderSplattingCUDACfg]): | |
| background_color: Float[Tensor, "3"] | |
| def __init__( | |
| self, | |
| cfg: DecoderSplattingCUDACfg, | |
| 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: DepthRenderingMode | None = None, | |
| ) -> DecoderOutput: | |
| b, v, _, _ = extrinsics.shape | |
| color = render_cuda( | |
| rearrange(extrinsics, "b v i j -> (b v) i j"), | |
| rearrange(intrinsics, "b v i j -> (b v) i j"), | |
| rearrange(near, "b v -> (b v)"), | |
| rearrange(far, "b v -> (b v)"), | |
| image_shape, | |
| repeat(self.background_color, "c -> (b v) c", b=b, v=v), | |
| repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v), | |
| repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v), | |
| repeat(gaussians.harmonics, "b g c d_sh -> (b v) g c d_sh", v=v), | |
| repeat(gaussians.opacities, "b g -> (b v) g", v=v), | |
| ) | |
| color = rearrange(color, "(b v) c h w -> b v c h w", b=b, v=v) | |
| return DecoderOutput( | |
| color, | |
| None | |
| if depth_mode is None | |
| else self.render_depth( | |
| gaussians, extrinsics, intrinsics, near, far, image_shape, depth_mode | |
| ), | |
| ) | |
| 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: DepthRenderingMode = "depth", | |
| ) -> Float[Tensor, "batch view height width"]: | |
| b, v, _, _ = extrinsics.shape | |
| result = render_depth_cuda( | |
| rearrange(extrinsics, "b v i j -> (b v) i j"), | |
| rearrange(intrinsics, "b v i j -> (b v) i j"), | |
| rearrange(near, "b v -> (b v)"), | |
| rearrange(far, "b v -> (b v)"), | |
| image_shape, | |
| repeat(gaussians.means, "b g xyz -> (b v) g xyz", v=v), | |
| repeat(gaussians.covariances, "b g i j -> (b v) g i j", v=v), | |
| repeat(gaussians.opacities, "b g -> (b v) g", v=v), | |
| mode=mode, | |
| ) | |
| return rearrange(result, "(b v) h w -> b v h w", b=b, v=v) | |