Add spherical harmonics support for Gaussian splats (new)

examples/01_scene/09_gaussian_splats.py

@@ -45,6 +45,9 @@ class SplatFile(TypedDict):
     """(N, 1). Range [0, 1]."""
     covariances: npt.NDArray[np.floating]
     """(N, 3, 3)."""
+    sh_coeffs: npt.NDArray[np.floating] | None
+    """(N, K, 3) spherical harmonics coefficients (including the DC term), or
+    None if the file only stores per-Gaussian colors."""
 
 
 def load_splat_file(splat_path: Path, center: bool = False) -> SplatFile:
@@ -88,6 +91,8 @@ def load_splat_file(splat_path: Path, center: bool = False) -> SplatFile:
         opacities=splat_uint8[:, 27:28] / 255.0,
         # Covariances should have shape (N, 3, 3).
         covariances=covariances,
+        # .splat files don't store spherical harmonics.
+        sh_coeffs=None,
     )
 
 
@@ -102,9 +107,26 @@ def load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile:
     positions = np.stack([v["x"], v["y"], v["z"]], axis=-1)
     scales = np.exp(np.stack([v["scale_0"], v["scale_1"], v["scale_2"]], axis=-1))
     wxyzs = np.stack([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], axis=1)
-    colors = 0.5 + SH_C0 * np.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1)
+    dc_terms = np.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1)
+    colors = 0.5 + SH_C0 * dc_terms
     opacities = 1.0 / (1.0 + np.exp(-v["opacity"][:, None]))
 
+    # Read higher-order spherical harmonics coefficients, if present. 3DGS
+    # checkpoints store them channel-major (all R terms, then G, then B); we
+    # want (N, K, 3) with the DC term first.
+    field_names = {prop.name for prop in v.properties}
+    num_rest = sum(1 for name in field_names if name.startswith("f_rest_"))
+    sh_coeffs = None
+    if num_rest > 0:
+        assert num_rest % 3 == 0
+        rest_per_channel = num_rest // 3
+        rest_terms = np.stack(
+            [v[f"f_rest_{i}"] for i in range(num_rest)], axis=1
+        ).reshape((-1, 3, rest_per_channel))
+        sh_coeffs = np.concatenate(
+            [dc_terms[:, :, None], rest_terms], axis=2
+        ).transpose(0, 2, 1)
+
     Rs = tf.SO3(wxyzs).as_matrix()
     covariances = np.einsum(
         "nij,njk,nlk->nil", Rs, np.eye(3)[None, :, :] * scales[:, None, :] ** 2, Rs
@@ -121,6 +143,7 @@ def load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile:
         rgbs=colors,
         opacities=opacities,
         covariances=covariances,
+        sh_coeffs=sh_coeffs,
     )
 
 
@@ -147,6 +170,7 @@ def main(
             rgbs=splat_data["rgbs"],
             opacities=splat_data["opacities"],
             covariances=splat_data["covariances"],
+            sh_coeffs=splat_data["sh_coeffs"],
         )
 
         remove_button = server.gui.add_button(f"Remove splat object {i}")

src/viser/_messages.py

@@ -2146,6 +2146,17 @@ class GaussianSplatsProps:
     - rgba (int32)
 
     Where cov1-6 are the upper-triangular terms of covariance matrices."""
+    sh_degree: int
+    """Spherical harmonics degree for view-dependent colors. 0 means no
+    spherical harmonics (colors are read from `buffer`)."""
+    sh_buffer: Optional[npt.NDArray[np.uint32]]
+    """Optional spherical harmonics coefficients, used when `sh_degree > 0`.
+
+    Each Gaussian gets `4 * ceil(3 * (sh_degree + 1)^2 / 8)` uint32 elements:
+    all `3 * (sh_degree + 1)^2` coefficients (including the DC term) as
+    float16, in coefficient-major order (c0.rgb, c1.rgb, ...), zero-padded to
+    a multiple of 8 float16s so each Gaussian spans a whole number of RGBA32UI
+    texels."""
     scale: Union[float, Tuple[float, float, float]] = 1.0
     """Scale of the Gaussian splats. A single float for uniform scaling or a
     tuple of (x, y, z) for per-axis scaling."""

src/viser/_scene_api.py

@@ -2169,6 +2169,7 @@ def add_gaussian_splats(
         rgbs: np.ndarray,
         opacities: np.ndarray,
         *,
+        sh_coeffs: np.ndarray | None = None,
         scale: float | tuple[float, float, float] = 1.0,
         wxyz: Tuple[float, float, float, float] | np.ndarray = (1.0, 0.0, 0.0, 0.0),
         position: Tuple[float, float, float] | np.ndarray = (0.0, 0.0, 0.0),
@@ -2185,6 +2186,13 @@ def add_gaussian_splats(
             covariances: Second moment for each Gaussian. (N, 3, 3).
             rgbs: Color for each Gaussian. (N, 3).
             opacities: Opacity for each Gaussian. (N, 1).
+            sh_coeffs: Optional spherical harmonics coefficients for
+                view-dependent colors, in the 3DGS (inria) convention. (N, K,
+                3), where K is 4, 9, or 16 for SH degrees 1, 2, and 3. The
+                first coefficient is the DC term (`f_dc` in standard 3DGS
+                checkpoints); the remainder are the higher-order terms
+                (`f_rest`), lowest order first. When provided, colors are
+                computed from the harmonics and `rgbs` is ignored.
             scale: Scale of the Gaussian splats. A single float for uniform
                 scaling or a tuple of (x, y, z) for per-axis scaling.
             wxyz: R_parent_local transformation.
@@ -2200,6 +2208,29 @@ def add_gaussian_splats(
         assert opacities.shape == (num_gaussians, 1)
         assert covariances.shape == (num_gaussians, 3, 3)
 
+        sh_degree = 0
+        sh_buffer = None
+        if sh_coeffs is not None:
+            degree_from_coeff_count = {4: 1, 9: 2, 16: 3}
+            assert (
+                sh_coeffs.ndim == 3
+                and sh_coeffs.shape[0] == num_gaussians
+                and sh_coeffs.shape[1] in degree_from_coeff_count
+                and sh_coeffs.shape[2] == 3
+            ), (
+                "sh_coeffs must have shape (N, K, 3) with K in (4, 9, 16),"
+                f" got {sh_coeffs.shape}"
+            )
+            sh_degree = degree_from_coeff_count[sh_coeffs.shape[1]]
+
+            # Pack coefficients as float16, zero-padded so each Gaussian spans
+            # a whole number of RGBA32UI texels (8 float16s each) client-side.
+            num_floats = sh_coeffs.shape[1] * 3
+            num_floats_padded = -(-num_floats // 8) * 8
+            sh_f16 = np.zeros((num_gaussians, num_floats_padded), dtype=np.float16)
+            sh_f16[:, :num_floats] = sh_coeffs.reshape(num_gaussians, num_floats)
+            sh_buffer = sh_f16.view(np.uint32)
+
         # Get upper-triangular terms of covariance matrix.
         cov_triu = covariances.reshape((-1, 9))[:, np.array([0, 1, 2, 4, 5, 8])]
         buffer = np.concatenate(
@@ -2225,6 +2256,8 @@ def add_gaussian_splats(
             props=_messages.GaussianSplatsProps(
                 buffer=buffer,
                 scale=scale,
+                sh_degree=sh_degree,
+                sh_buffer=sh_buffer,
             ),
         )
         node_handle = GaussianSplatHandle._make(

src/viser/_scene_handles.py

@@ -999,6 +999,12 @@ class GaussianSplatHandle(
         - [3]: reserved for renderer
         - [4:7]: covariance upper-triangular (6x float16)
         - [7]: RGBA (4x uint8)
+
+    When `sh_degree > 0`, view-dependent colors are computed from the
+    float16 spherical harmonics coefficients in `sh_buffer` instead of the
+    RGB values in `buffer`. The sub-property setters below (`centers`,
+    `rgbs`, ...) update `buffer` only; if they change the number of
+    Gaussians, a stale `sh_buffer` is ignored by the renderer.
     """
 
     def _ensure_buffer_size(self, num_gaussians: int) -> None:

src/viser/client/src/SceneTree.tsx

@@ -580,6 +580,8 @@ function createObjectFactory(
             <group scale={normalizeScale(message.props.scale)}>
               <SplatObject
                 buffer={message.props.buffer}
+                shBuffer={message.props.sh_buffer}
+                shDegree={message.props.sh_degree}
                 sceneNodeName={message.name}
               />
             </group>