diff --git a/nerfstudio/models/splatfacto.py b/nerfstudio/models/splatfacto.py index 6eafd8013..6ce5c164f 100644 --- a/nerfstudio/models/splatfacto.py +++ b/nerfstudio/models/splatfacto.py @@ -748,7 +748,7 @@ class SplatfactoModel(Model): # 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 @@ class SplatfactoModel(Model): 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 @@ class SplatfactoModel(Model): 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 diff --git a/pyproject.toml b/pyproject.toml index 6ee1f42d0..c4ea5fc15 100644 --- a/pyproject.toml +++ b/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"