Spaces:
Running on Zero
Running on Zero
File size: 5,006 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
import torch
from einops import rearrange
from jaxtyping import Float
from torch import Tensor, nn
import torch.nn.functional as F
from ....geometry.projection import get_world_rays
from ....misc.sh_rotation import rotate_sh
from .gaussians import build_covariance
@dataclass
class Gaussians:
means: Float[Tensor, "*batch 3"]
covariances: Float[Tensor, "*batch 3 3"]
scales: Float[Tensor, "*batch 3"]
rotations: Float[Tensor, "*batch 4"]
harmonics: Float[Tensor, "*batch 3 _"]
opacities: Float[Tensor, " *batch"]
rotations_unnorm: Float[Tensor, "*batch 4"]
@dataclass
class GaussianAdapterCfg:
gaussian_scale_min: float
gaussian_scale_max: float
sh_degree: int
exp_scale_bias: float
no_rotate_sh: bool
no_sh_mask: bool
class GaussianAdapter(nn.Module):
cfg: GaussianAdapterCfg
def __init__(self, cfg: GaussianAdapterCfg):
super().__init__()
self.cfg = cfg
# Create a mask for the spherical harmonics coefficients. This ensures that at
# initialization, the coefficients are biased towards having a large DC
# component and small view-dependent components.
self.register_buffer(
"sh_mask",
torch.ones((self.d_sh,), dtype=torch.float32),
persistent=False,
)
for degree in range(1, self.cfg.sh_degree + 1):
self.sh_mask[degree**2 : (degree + 1) ** 2] = 0.1 * 0.25**degree
def forward(
self,
extrinsics: Float[Tensor, "*#batch 4 4"],
intrinsics: Float[Tensor, "*#batch 3 3"] | None,
coordinates: Float[Tensor, "*#batch 2"],
depths: Float[Tensor, "*#batch"] | None,
opacities: Float[Tensor, "*#batch"],
raw_gaussians: Float[Tensor, "*#batch _"],
image_shape: tuple[int, int],
eps: float = 1e-8,
point_cloud: Float[Tensor, "*#batch 3"] | None = None,
input_images: Tensor | None = None,
) -> Gaussians:
# Gaussian geometry is consumed by CUDA/gsplat rasterizers, so keep the
# position, covariance, and SH rotation math in FP32 under mixed precision.
with torch.autocast(device_type=raw_gaussians.device.type, enabled=False):
extrinsics = extrinsics.float()
intrinsics = intrinsics.float() if intrinsics is not None else None
coordinates = coordinates.float()
depths = depths.float() if depths is not None else None
opacities = opacities.float()
raw_gaussians = raw_gaussians.float()
scales, rotations_unnorm, sh = raw_gaussians.split(
(3, 4, 3 * self.d_sh), dim=-1
)
scales = torch.clamp(
F.softplus(scales - self.cfg.exp_scale_bias),
min=self.cfg.gaussian_scale_min,
max=self.cfg.gaussian_scale_max,
)
# Normalize the quaternion features to yield a valid quaternion.
rotations = rotations_unnorm / (
rotations_unnorm.norm(dim=-1, keepdim=True) + eps
)
# Split RGB spherical-harmonic coefficients into the expected layout.
sh = rearrange(sh, "... (xyz d_sh) -> ... xyz d_sh", xyz=3)
sh = sh.broadcast_to((*opacities.shape, 3, self.d_sh))
if self.cfg.no_sh_mask:
sh = sh.clone()
else:
sh = sh * self.sh_mask
if input_images is not None:
# [B, V, H*W, 1, 1, 3]
imgs = rearrange(
input_images.float(), "b v c h w -> b v (h w) () () c"
)
# init sh with input images
sh[..., 0] = sh[..., 0] + RGB2SH(imgs)
# Create world-space covariance matrices.
covariances = build_covariance(scales, rotations)
c2w_rotations = extrinsics[..., :3, :3]
covariances = (
c2w_rotations @ covariances @ c2w_rotations.transpose(-1, -2)
)
# Compute Gaussian means.
origins, directions = get_world_rays(coordinates, extrinsics, intrinsics)
means = origins + directions * depths[..., None]
harmonics = (
sh
if self.cfg.no_rotate_sh
else rotate_sh(sh, c2w_rotations[..., None, :, :])
)
return Gaussians(
means=means,
covariances=covariances,
harmonics=harmonics,
opacities=opacities,
scales=scales,
rotations=rotations.broadcast_to((*scales.shape[:-1], 4)),
rotations_unnorm=rotations_unnorm.broadcast_to((*scales.shape[:-1], 4)),
)
@property
def d_sh(self) -> int:
return (self.cfg.sh_degree + 1) ** 2
@property
def d_in(self) -> int:
return 7 + 3 * self.d_sh
def RGB2SH(rgb):
C0 = 0.28209479177387814
return (rgb - 0.5) / C0
|