param tweaks, dont generate raybundle in viewer

nerfstudio/models/gaussian_splatting.py

@@ -44,6 +44,91 @@ from diff_rast.rasterize import RasterizeGaussians
 from diff_rast.project_gaussians import ProjectGaussians
 from diff_rast.sh import SphericalHarmonics, num_sh_bases
 
+def eval_sh(deg, sh, dirs):
+    """
+    Evaluate spherical harmonics at unit directions
+    using hardcoded SH polynomials.
+    Works with torch/np/jnp.
+    ... Can be 0 or more batch dimensions.
+    Args:
+        deg: int SH deg. Currently, 0-3 supported
+        sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
+        dirs: jnp.ndarray unit directions [..., 3]
+    Returns:
+        [..., C]
+    """
+    C0 = 0.28209479177387814
+    C1 = 0.4886025119029199
+    C2 = [1.0925484305920792, -1.0925484305920792, 0.31539156525252005, -1.0925484305920792, 0.5462742152960396]
+    C3 = [
+        -0.5900435899266435,
+        2.890611442640554,
+        -0.4570457994644658,
+        0.3731763325901154,
+        -0.4570457994644658,
+        1.445305721320277,
+        -0.5900435899266435,
+    ]
+    C4 = [
+        2.5033429417967046,
+        -1.7701307697799304,
+        0.9461746957575601,
+        -0.6690465435572892,
+        0.10578554691520431,
+        -0.6690465435572892,
+        0.47308734787878004,
+        -1.7701307697799304,
+        0.6258357354491761,
+    ]
+
+    assert deg <= 4 and deg >= 0
+    coeff = (deg + 1) ** 2
+    assert sh.shape[-1] >= coeff
+
+    result = C0 * sh[..., 0]
+    if deg > 0:
+        x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
+        result = result - C1 * y * sh[..., 1] + C1 * z * sh[..., 2] - C1 * x * sh[..., 3]
+
+        if deg > 1:
+            xx, yy, zz = x * x, y * y, z * z
+            xy, yz, xz = x * y, y * z, x * z
+            result = (
+                result
+                + C2[0] * xy * sh[..., 4]
+                + C2[1] * yz * sh[..., 5]
+                + C2[2] * (2.0 * zz - xx - yy) * sh[..., 6]
+                + C2[3] * xz * sh[..., 7]
+                + C2[4] * (xx - yy) * sh[..., 8]
+            )
+
+            if deg > 2:
+                result = (
+                    result
+                    + C3[0] * y * (3 * xx - yy) * sh[..., 9]
+                    + C3[1] * xy * z * sh[..., 10]
+                    + C3[2] * y * (4 * zz - xx - yy) * sh[..., 11]
+                    + C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12]
+                    + C3[4] * x * (4 * zz - xx - yy) * sh[..., 13]
+                    + C3[5] * z * (xx - yy) * sh[..., 14]
+                    + C3[6] * x * (xx - 3 * yy) * sh[..., 15]
+                )
+
+                if deg > 3:
+                    result = (
+                        result
+                        + C4[0] * xy * (xx - yy) * sh[..., 16]
+                        + C4[1] * yz * (3 * xx - yy) * sh[..., 17]
+                        + C4[2] * xy * (7 * zz - 1) * sh[..., 18]
+                        + C4[3] * yz * (7 * zz - 3) * sh[..., 19]
+                        + C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20]
+                        + C4[5] * xz * (7 * zz - 3) * sh[..., 21]
+                        + C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22]
+                        + C4[7] * xz * (xx - 3 * yy) * sh[..., 23]
+                        + C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24]
+                    )
+    return result
+
 def random_quat_tensor(N, **kwargs):
     u = torch.rand(N, **kwargs)
     v = torch.rand(N, **kwargs)
@@ -103,9 +188,9 @@ class GaussianSplattingModelConfig(ModelConfig):
     """threshold of positional gradient norm for densifying gaussians"""
     densify_size_thresh: float = 0.01
     """below this size, gaussians are *duplicated*, otherwise split"""
-    sh_degree_interval: int = 1000
+    sh_degree_interval: int = 500
     """every n intervals turn on another sh degree"""
-    max_screen_size: int = 300
+    max_screen_size: int = 150
     """maximum screen size of a gaussian, in pixels"""
     random_init: bool = False
     """whether to initialize the positions uniformly randomly (not SFM points)"""
@@ -324,8 +409,8 @@ class GaussianSplattingModel(Model):
         # cull transparent ones
         culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze()
         if self.step > self.config.refine_every * self.config.reset_alpha_every:
-            # cull huge ones that are in the center of the scene (a 2-unit box is twice as large as all the cameras)
-            toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() & (self.means.amax(dim=-1) < 2)
+            # cull huge ones
+            toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze()
             culls = culls | toobigs
             # cull big screen space
             culls = culls | (self.max_2Dsize > self.config.max_screen_size).squeeze()
@@ -507,7 +592,7 @@ class GaussianSplattingModel(Model):
             viewdirs = means_crop[rend_mask].detach() - camera.camera_to_worlds[..., :3, 3]  # (N, 3)
             viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True)
             n = min(self.step // self.config.sh_degree_interval,self.config.sh_degree)
-            # rgbs = eval_sh(n, self.get_colors, viewdirs)[crop_ids]
+            # rgbs = eval_sh(n, colors_crop[rend_mask,:,:], viewdirs)
             n_coeffs = num_sh_bases(n)
             coeffs = colors_crop.permute(0,2,1)[rend_mask,:n_coeffs,:]
             rgbs = SphericalHarmonics.apply(n, viewdirs,coeffs)
@@ -528,19 +613,19 @@ class GaussianSplattingModel(Model):
             W,
             background,
         )
-        depth_im = RasterizeGaussians.apply(
-            self.xys[rend_mask],
-            depths[rend_mask],
-            self.radii[rend_mask],
-            conics[rend_mask],
-            num_tiles_hit[rend_mask],
-            depths[:,None],
-            torch.sigmoid(opacities_crop[rend_mask]),
-            H,
-            W,
-            torch.ones(1,device=self.device)*10,
-        )[...,0:1]
-        return {"rgb": rgb[..., :3],'depth':depth_im}
+        # depth_im = RasterizeGaussians.apply(
+        #     self.xys[rend_mask],
+        #     depths[rend_mask],
+        #     self.radii[rend_mask],
+        #     conics[rend_mask],
+        #     num_tiles_hit[rend_mask],
+        #     depths[rend_mask,None],
+        #     torch.sigmoid(opacities_crop[rend_mask]),
+        #     H,
+        #     W,
+        #     torch.ones(1,device=self.device)*10,
+        # )[...,0:1]
+        return {"rgb": rgb}#,'depth':depth_im}
 
     def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]:
         """Compute and returns metrics.

nerfstudio/viewer_beta/render_state_machine.py

@@ -26,7 +26,7 @@ from nerfstudio.utils import colormaps, writer
 from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
 from nerfstudio.viewer.server import viewer_utils
 from nerfstudio.viewer_beta.utils import CameraState, get_camera
-
+from nerfstudio.models.gaussian_splatting import GaussianSplattingModel
 if TYPE_CHECKING:
     from nerfstudio.viewer_beta.viewer import Viewer
 
@@ -126,7 +126,10 @@ class RenderStateMachine(threading.Thread):
 
         with TimeWriter(None, None, write=False) as vis_t:
             with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext():
-                camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb)
+                if isinstance(self.viewer.get_model(),GaussianSplattingModel):
+                    camera_ray_bundle=None
+                else:
+                    camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb)
                 self.viewer.get_model().set_crop(obb)
                 self.viewer.get_model().eval()
                 step = self.viewer.step