Add alpha rendering to splatfacto, fix depth bug (nerfstudio-project#…

…2824)

* Add alpha rendering to splatfacto, output as accumulation

nerfstudio/models/splatfacto.py

@@ -748,7 +748,7 @@ def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
         # rescale the camera back to original dimensions
         camera.rescale_output_resolution(camera_downscale)
         assert (num_tiles_hit > 0).any()  # type: ignore
-        rgb = rasterize_gaussians(  # type: ignore
+        rgb, alpha = rasterize_gaussians(  # type: ignore
             self.xys,
             depths,
             self.radii,
@@ -759,7 +759,9 @@ def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
             H,
             W,
             background=background,
+            return_alpha=True,
         )  # type: ignore
+        alpha = alpha[..., None]
         rgb = torch.clamp(rgb, max=1.0)  # type: ignore
         depth_im = None
         if not self.training:
@@ -773,10 +775,12 @@ def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
                 torch.sigmoid(opacities_crop),
                 H,
                 W,
-                background=torch.ones(3, device=self.device) * 10,
+                background=torch.zeros(3, device=self.device),
             )[..., 0:1]  # type: ignore
+            depth_im[alpha > 0] = depth_im[alpha > 0] / alpha[alpha > 0]
+            depth_im[alpha == 0] = 1000
 
-        return {"rgb": rgb, "depth": depth_im}  # type: ignore
+        return {"rgb": rgb, "depth": depth_im, "accumulation": alpha}  # type: ignore
 
     def get_gt_img(self, image: torch.Tensor):
         """Compute groundtruth image with iteration dependent downscale factor for evaluation purpose

pyproject.toml

@@ -63,7 +63,7 @@ dependencies = [
     "xatlas",
     "trimesh>=3.20.2",
     "timm==0.6.7",
-    "gsplat==0.1.2.1",
+    "gsplat==0.1.3",
     "pytorch-msssim",
     "pathos",
     "packaging"