Migrate to python functions for gsplat (#2778)

* switch to pythonic interface for gsplat

nerfstudio/models/gaussian_splatting.py

@@ -26,10 +26,9 @@ from typing import Dict, List, Optional, Tuple, Type, Union
 import numpy as np
 import torch
 from gsplat._torch_impl import quat_to_rotmat
-from gsplat.compute_cumulative_intersects import compute_cumulative_intersects
-from gsplat.project_gaussians import ProjectGaussians
-from gsplat.rasterize import RasterizeGaussians
-from gsplat.sh import SphericalHarmonics, num_sh_bases
+from gsplat.project_gaussians import project_gaussians
+from gsplat.rasterize import rasterize_gaussians
+from gsplat.sh import num_sh_bases, spherical_harmonics
 from pytorch_msssim import SSIM
 from torch.nn import Parameter
 
@@ -324,7 +323,8 @@ class GaussianSplattingModel(Model):
         with torch.no_grad():
             # keep track of a moving average of grad norms
             visible_mask = (self.radii > 0).flatten()
-            grads = self.xys.grad.detach().norm(dim=-1)  # TODO fill in
+            assert self.xys.grad is not None
+            grads = self.xys.grad.detach().norm(dim=-1)
             # print(f"grad norm min {grads.min().item()} max {grads.max().item()} mean {grads.mean().item()} size {grads.shape}")
             if self.xys_grad_norm is None:
                 self.xys_grad_norm = grads
@@ -629,13 +629,13 @@ class GaussianSplattingModel(Model):
         cy = camera.cy.item()
         fovx = 2 * math.atan(camera.width / (2 * camera.fx))
         fovy = 2 * math.atan(camera.height / (2 * camera.fy))
-        W, H = camera.width.item(), camera.height.item()
+        W, H = int(camera.width.item()), int(camera.height.item())
         self.last_size = (H, W)
         projmat = projection_matrix(0.001, 1000, fovx, fovy, device=self.device)
         BLOCK_X, BLOCK_Y = 16, 16
         tile_bounds = (
-            (W + BLOCK_X - 1) // BLOCK_X,
-            (H + BLOCK_Y - 1) // BLOCK_Y,
+            int((W + BLOCK_X - 1) // BLOCK_X),
+            int((H + BLOCK_Y - 1) // BLOCK_Y),
             1,
         )
 
@@ -656,7 +656,7 @@ class GaussianSplattingModel(Model):
 
         colors_crop = torch.cat((features_dc_crop[:, None, :], features_rest_crop), dim=1)
 
-        self.xys, depths, self.radii, conics, num_tiles_hit, cov3d = ProjectGaussians.apply(
+        self.xys, depths, self.radii, conics, num_tiles_hit, cov3d = project_gaussians(  # type: ignore
             means_crop,
             torch.exp(scales_crop),
             1,
@@ -682,7 +682,7 @@ class GaussianSplattingModel(Model):
             viewdirs = means_crop.detach() - camera.camera_to_worlds.detach()[..., :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 = SphericalHarmonics.apply(n, viewdirs, colors_crop)
+            rgbs = spherical_harmonics(n, viewdirs, colors_crop)
             rgbs = torch.clamp(rgbs + 0.5, min=0.0)  # type: ignore
         else:
             rgbs = torch.sigmoid(colors_crop[:, 0, :])
@@ -690,36 +690,34 @@ class GaussianSplattingModel(Model):
         # rescale the camera back to original dimensions
         camera.rescale_output_resolution(camera_downscale)
 
-        # avoid empty rasterization
-        num_intersects, _ = compute_cumulative_intersects(self.xys.size(0), num_tiles_hit)
-        assert num_intersects > 0
+        assert (num_tiles_hit > 0).any()  # type: ignore
 
-        rgb = RasterizeGaussians.apply(
+        rgb = rasterize_gaussians(  # type: ignore
             self.xys,
             depths,
             self.radii,
             conics,
-            num_tiles_hit,
+            num_tiles_hit,  # type: ignore
             rgbs,
             torch.sigmoid(opacities_crop),
             H,
             W,
-            background,
+            background=background,
         )  # type: ignore
         rgb = torch.clamp(rgb, max=1.0)  # type: ignore
         depth_im = None
         if not self.training:
-            depth_im = RasterizeGaussians.apply(  # type: ignore
+            depth_im = rasterize_gaussians(  # type: ignore
                 self.xys,
                 depths,
                 self.radii,
                 conics,
-                num_tiles_hit,
+                num_tiles_hit,  # type: ignore
                 depths[:, None].repeat(1, 3),
                 torch.sigmoid(opacities_crop),
                 H,
                 W,
-                torch.ones(3, device=self.device) * 10,
+                background=torch.ones(3, device=self.device) * 10,
             )[..., 0:1]  # type: ignore
 
         return {"rgb": rgb, "depth": depth_im}  # type: ignore

pyproject.toml

@@ -62,7 +62,7 @@ dependencies = [
     "xatlas",
     "trimesh>=3.20.2",
     "timm==0.6.7",
-    "gsplat==0.1.0",
+    "gsplat==0.1.2.1",
     "pytorch-msssim",
     "pathos"
 ]