updating background for different rendering modes [D, ED, RGB+D, RGB+… (

#240) * updating background for different rendering modes [D, ED, RGB+D, RGB+ED] to be compatible for rasterization * updating formatting to black 22.3.0
nerfstudio-project · Jun 27, 2024 · 880425c · 880425c
1 parent 2aa70c2
commit 880425c
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/gsplat/rendering.py b/gsplat/rendering.py
@@ -317,8 +317,14 @@ def rasterization(
     # Rasterize to pixels
     if render_mode in ["RGB+D", "RGB+ED"]:
         colors = torch.cat((colors, depths[..., None]), dim=-1)
+        if backgrounds is not None:
+            backgrounds = torch.cat(
+                [backgrounds, torch.zeros(C, 1, device=backgrounds.device)], dim=-1
+            )
     elif render_mode in ["D", "ED"]:
         colors = depths[..., None]
+        if backgrounds is not None:
+            backgrounds = torch.zeros(C, 1, device=backgrounds.device)
     else:  # RGB
         pass
     if colors.shape[-1] > channel_chunk: