Spherical harmonics support

examples/experimental/gaussian_splats.py

@@ -82,8 +82,30 @@ def load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile:
     positions = onp.stack([v["x"], v["y"], v["z"]], axis=-1)
     scales = onp.exp(onp.stack([v["scale_0"], v["scale_1"], v["scale_2"]], axis=-1))
     wxyzs = onp.stack([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], axis=1)
-    colors = 0.5 + SH_C0 * onp.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1)
+    colors = 0.5 + SH_C0 * onp.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1) 
+    # print(v["f_dc_0"].shape) # prints (numGaussians)
+    # print(colors.shape) # prints (numGaussians, 3)
     opacities = 1.0 / (1.0 + onp.exp(-v["opacity"][:, None]))
+
+    # Load all zero order SH coefficients
+    dc_terms = onp.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1)
+
+    # Load higher order SH coefficients (f_rest_0, ... f_rest_44), which are either level 1 or higher
+    # Note: .ply file supports maximum SH degree of 3, R = f_rest_0, ... f_rest_14; G = f_rest_15, ... f_rest_29
+    rest_terms = []
+    i = 0
+    #while f"f_rest_{i}" in v:
+    while i < 15:
+        rest_terms.append(v[f"f_rest_{i}"]) # has shape (numGaussians, )
+        rest_terms.append(v[f"f_rest_{15 + i}"])
+        rest_terms.append(v[f"f_rest_{30 + i}"])
+        i += 1
+    # while f"f_rest_{i}" in v:
+    #     rest_terms.append(v[f"f_rest_{i}"])
+    #     i += 1
+    if len(rest_terms) > 0: # if we do have higher than zero order SH, we will process them and add them here.
+        sh_coeffs = onp.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]] + rest_terms, axis=1)
+    sh_degree = int(onp.sqrt(sh_coeffs.shape[1] // 3) - 1)
 
     Rs = tf.SO3(wxyzs).as_matrix()
     covariances = onp.einsum(
@@ -93,19 +115,42 @@ def load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile:
         positions -= onp.mean(positions, axis=0, keepdims=True)
 
     num_gaussians = len(v)
+
+    # print(sh_coeffs.shape) # prints (447703, 48)
+    # print(v["x"].shape) # prints (447703,)
+    # print(positions.shape)
+    # print(colors.shape)
+    # print(covariances.shape)
+
     print(
         f"PLY file with {num_gaussians=} loaded in {time.time() - start_time} seconds"
     )
+    print(onp.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1)[0 , :])
+    print(sh_coeffs[0, :]) # first gaussian, all 48 coefficients
+    # print(sh_coeffs[0, :]) # next 3 SH coefficients that are the 1st order
+
+    # return {
+    #     "centers": positions[0:1, :],
+    #     "rgbs": colors[0:1, :],
+    #     "opacities": opacities[0:1, :],
+    #     "covariances": 10000*covariances[0:1, :],
+    #     "sh_degree": sh_degree,
+    #     "sh_coeffs": sh_coeffs[0:1, :],
+    # }
+
     return {
         "centers": positions,
         "rgbs": colors,
         "opacities": opacities,
         "covariances": covariances,
+        "sh_degree": sh_degree,
+        "sh_coeffs": sh_coeffs,
     }
 
 
 def main(splat_paths: tuple[Path, ...]) -> None:
     server = viser.ViserServer()
+    print(server.request_share_url())
     server.gui.configure_theme(dark_mode=True)
     gui_reset_up = server.gui.add_button(
         "Reset up direction",
@@ -135,6 +180,8 @@ def _(event: viser.GuiEvent) -> None:
             rgbs=splat_data["rgbs"],
             opacities=splat_data["opacities"],
             covariances=splat_data["covariances"],
+            sh_degree=splat_data["sh_degree"],
+            sh_coeffs=splat_data["sh_coeffs"],
         )
 
         remove_button = server.gui.add_button(f"Remove splat object {i}")
@@ -150,3 +197,7 @@ def _(_, gs_handle=gs_handle, remove_button=remove_button) -> None:
 
 if __name__ == "__main__":
     tyro.cli(main)
+
+
+
+print("yapyap")

src/viser/_messages.py

@@ -761,6 +761,14 @@ class GaussianSplatsMessage(Message):
     - cov5 (f16), cov6 (f16)
     - rgba (int32)
     Where cov1-6 are the upper triangular elements of the covariance matrix."""
+
+    sh_buffer: onpt.NDArray[onp.uint32]
+    """The spherical harmonic buffer contains:
+    - 
+    -
+    -
+    """
+
 
 
 @dataclasses.dataclass

src/viser/_scene_api.py

@@ -942,6 +942,8 @@ def _add_gaussian_splats(
         wxyz: Tuple[float, float, float, float] | onp.ndarray = (1.0, 0.0, 0.0, 0.0),
         position: Tuple[float, float, float] | onp.ndarray = (0.0, 0.0, 0.0),
         visible: bool = True,
+        sh_degree: int = 0,
+        sh_coeffs: onp.ndarray = None,
     ) -> GaussianSplatHandle:
         """Add a model to render using Gaussian Splatting.
 
@@ -990,12 +992,24 @@ def _add_gaussian_splats(
             ],
             axis=-1,
         ).view(onp.uint32)
-        assert buffer.shape == (num_gaussians, 8)
+
+        assert buffer.shape == (num_gaussians, 8) 
+
+        # We have 48 float32 spherical coeffecients per gaussian. 
+        # However, by converting them to float16 we now have 48 float16 values per gaussian
+        # This means, each cell of sh_buffer contains 2 spherical coefficients because each cell is 32bits
+
+        # - (768 bits): spherical harmonics
+        print("sh_coeffs.shape", sh_coeffs.shape)
+        sh_buffer = (sh_coeffs.astype(onp.float16)).view(onp.uint32)
+        print("sh_buffer.shape", sh_buffer.shape) # has shape (num_gaussians, 24), each cell contains 2 spherical coeff.
+        # print(sh_buffer)
 
         self._websock_interface.queue_message(
             _messages.GaussianSplatsMessage(
                 name=name,
                 buffer=buffer,
+                sh_buffer=sh_buffer,
             )
         )
         node_handle = GaussianSplatHandle._make(self, name, wxyz, position, visible)

src/viser/client/.yarnrc.yml

@@ -0,0 +1 @@
+nodeLinker: node-modules

src/viser/client/src/MessageHandler.tsx

@@ -996,6 +996,14 @@ function useMessageHandler() {
                     ),
                   )
                 }
+                sh_buffer={
+                  new Uint32Array(
+                    message.sh_buffer.buffer.slice(
+                      message.sh_buffer.byteOffset,
+                      message.sh_buffer.byteOffset + message.sh_buffer.byteLength,
+                    ),
+                  )
+                }
               />
             );
           }),