Add default depth output when no gaussians visible (#2874)

* Add default depth output when no gaussians visible

* Add accumulation to default output return

* Change to using new_ones/zeros
This commit is contained in:
Chung Min Kim authored and GitHub committed 2024-02-07 00:21:21 +00:00
1 parent 6127ec0f72
commit dc53517565
1 file changed
+8 -2
+8 -2
View File
@@ -669,7 +669,10 @@ class SplatfactoModel(Model):
if self.crop_box is not None and not self.training:
crop_ids = self.crop_box.within(self.means).squeeze()
if crop_ids.sum() == 0:
return {"rgb": background.repeat(int(camera.height.item()), int(camera.width.item()), 1)}
rgb = background.repeat(int(camera.height.item()), int(camera.width.item()), 1)
depth = background.new_ones(*rgb.shape[:2], 1) * 10
accumulation = background.new_zeros(*rgb.shape[:2], 1)
return {"rgb": rgb, "depth": depth, "accumulation": accumulation}
else:
crop_ids = None
camera_downscale = self._get_downscale_factor()
@@ -734,7 +737,10 @@ class SplatfactoModel(Model):
tile_bounds,
) # type: ignore
if (self.radii).sum() == 0:
return {"rgb": background.repeat(int(camera.height.item()), int(camera.width.item()), 1)}
rgb = background.repeat(int(camera.height.item()), int(camera.width.item()), 1)
depth = background.new_ones(*rgb.shape[:2], 1) * 10
accumulation = background.new_zeros(*rgb.shape[:2], 1)
return {"rgb": rgb, "depth": depth, "accumulation": accumulation}
# Important to allow xys grads to populate properly
if self.training: