Spaces:
Running on Zero
Running on Zero
File size: 4,667 Bytes
7405e85 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | 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,
)
@dataclass
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,
)
|