add return_alpha keyword

nerfstudio-project · Jan 14, 2024 · 895ea08 · 895ea08
1 parent f72dd2a
commit 895ea08
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/gsplat/nd_rasterize.py b/gsplat/nd_rasterize.py
diff --git a/gsplat/rasterize.py b/gsplat/rasterize.py
@@ -22,6 +22,7 @@ def rasterize_gaussians(
     img_height: int,
     img_width: int,
     background: Optional[Float[Tensor, "channels"]] = None,
+    return_alpha: Optional[bool] = False,
 ) -> Tensor:
     """Rasterizes 2D gaussians by sorting and binning gaussian intersections for each tile and returns an N-dimensional output using alpha-compositing.
 
@@ -75,6 +76,7 @@ def rasterize_gaussians(
         img_height,
         img_width,
         background.contiguous(),
+        return_alpha,
     )
 
 
@@ -94,6 +96,7 @@ def forward(
         img_height: int,
         img_width: int,
         background: Optional[Float[Tensor, "channels"]] = None,
+        return_alpha: Optional[bool] = False,
     ) -> Tensor:
         num_points = xys.size(0)
         BLOCK_X, BLOCK_Y = 16, 16
@@ -147,9 +150,12 @@ def forward(
             final_Ts,
             final_idx,
         )
-        out_alpha = 1 - final_Ts
-
-        return out_img, out_alpha
+
+        if return_alpha:
+            out_alpha = 1 - final_Ts
+            return out_img, out_alpha
+        else:
+            return out_img
 
     @staticmethod
     def backward(ctx, v_out_img, v_out_alpha=None):
@@ -158,7 +164,7 @@ def backward(ctx, v_out_img, v_out_alpha=None):
 
         if v_out_alpha is None:
             v_out_alpha = torch.zeros_like(v_out_img[..., 0])
-
+        print(v_out_alpha)
         (
             gaussian_ids_sorted,
             tile_bins,
@@ -202,4 +208,5 @@ def backward(ctx, v_out_img, v_out_alpha=None):
             None,  # img_height
             None,  # img_width
             None,  # background
+            None,  # return_alpha
         )