plumb to python and add pytests with backgrounds

Signed-off-by: Mark Harris <mharris@nvidia.com>

Author: harrism

fvdb/__init__.py

@@ -108,6 +108,7 @@ def gaussian_render_jagged(
     render_depth_channel: bool = False,
     return_debug_info: bool = False,
     ortho: bool = False,
+    backgrounds: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
     return _gaussian_render_jagged_cpp(
         means=means._impl,
@@ -129,6 +130,7 @@ def gaussian_render_jagged(
         render_depth_channel=render_depth_channel,
         return_debug_info=return_debug_info,
         ortho=ortho,
+        backgrounds=backgrounds,
     )
 
 

fvdb/__init__.pyi

@@ -90,7 +90,8 @@ def gaussian_render_jagged(
     return_debug_info: bool = False,
     render_depth_only: bool = False,
     ortho: bool = False,
-) -> JaggedTensor: ...
+    backgrounds: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: ...
 @overload
 def jcat(grid_batches: Sequence[GridBatch]) -> GridBatch: ...
 @overload

fvdb/_fvdb_cpp.pyi

@@ -3,7 +3,7 @@
 #
 import typing
 from enum import Enum
-from typing import ClassVar, overload
+from typing import ClassVar, Optional, overload
 
 import torch
 
@@ -165,6 +165,7 @@ class GaussianSplat3d:
         min_radius_2d: float = ...,
         eps_2d: float = ...,
         antialias: bool = ...,
+        backgrounds: Optional[torch.Tensor] = ...,
     ) -> tuple[torch.Tensor, torch.Tensor]: ...
     def render_from_projected_gaussians(
         self,
@@ -174,6 +175,7 @@ class GaussianSplat3d:
         crop_origin_w: int = ...,
         crop_origin_h: int = ...,
         tile_size: int = ...,
+        backgrounds: Optional[torch.Tensor] = ...,
     ) -> tuple[torch.Tensor, torch.Tensor]: ...
     def render_images(
         self,
@@ -189,6 +191,7 @@ class GaussianSplat3d:
         min_radius_2d: float = ...,
         eps_2d: float = ...,
         antialias: bool = ...,
+        backgrounds: Optional[torch.Tensor] = ...,
     ) -> tuple[torch.Tensor, torch.Tensor]: ...
     def render_images_and_depths(
         self,
@@ -204,6 +207,7 @@ class GaussianSplat3d:
         min_radius_2d: float = ...,
         eps_2d: float = ...,
         antialias: bool = ...,
+        backgrounds: Optional[torch.Tensor] = ...,
     ) -> tuple[torch.Tensor, torch.Tensor]: ...
     def render_num_contributing_gaussians(
         self,
@@ -1037,7 +1041,29 @@ class config:
     pedantic_error_checking: ClassVar[bool] = ...
     def __init__(self, *args, **kwargs) -> None: ...
 
-def gaussian_render_jagged(*args, **kwargs): ...
+def gaussian_render_jagged(
+    means: JaggedTensor,
+    quats: JaggedTensor,
+    scales: JaggedTensor,
+    opacities: JaggedTensor,
+    sh_coeffs: JaggedTensor,
+    viewmats: JaggedTensor,
+    Ks: JaggedTensor,
+    image_width: int,
+    image_height: int,
+    near_plane: float = ...,
+    far_plane: float = ...,
+    sh_degree_to_use: int = ...,
+    tile_size: int = ...,
+    radius_clip: float = ...,
+    eps2d: float = ...,
+    antialias: bool = ...,
+    render_depth_channel: bool = ...,
+    return_debug_info: bool = ...,
+    render_depth_only: bool = ...,
+    ortho: bool = ...,
+    backgrounds: Optional[torch.Tensor] = ...,
+) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]: ...
 @overload
 def gridbatch_from_dense(
     num_grids: int,

fvdb/gaussian_splatting.py

@@ -1529,6 +1529,7 @@ def render_from_projected_gaussians(
         crop_origin_w: int = -1,
         crop_origin_h: int = -1,
         tile_size: int = 16,
+        backgrounds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Render a set of images from Gaussian splats that have already been projected onto image planes
@@ -1613,6 +1614,7 @@ def render_from_projected_gaussians(
             crop_origin_w=crop_origin_w,
             crop_origin_h=crop_origin_h,
             tile_size=tile_size,
+            backgrounds=backgrounds,
         )
 
     def render_depths(
@@ -1628,6 +1630,7 @@ def render_depths(
         min_radius_2d: float = 0.3,
         eps_2d: float = 0.3,
         antialias: bool = False,
+        backgrounds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Render ``C`` depth maps from this :class:`GaussianSplat3d` from ``C`` camera views.
@@ -1693,6 +1696,7 @@ def render_depths(
             min_radius_2d=min_radius_2d,
             eps_2d=eps_2d,
             antialias=antialias,
+            backgrounds=backgrounds,
         )
 
     def render_images(
@@ -1709,6 +1713,7 @@ def render_images(
         min_radius_2d: float = 0.0,
         eps_2d: float = 0.3,
         antialias: bool = False,
+        backgrounds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Render ``C`` multi-channel images (see :attr:`num_channels`) from this :class:`GaussianSplat3d` from ``C`` camera views.
@@ -1775,6 +1780,7 @@ def render_images(
             min_radius_2d=min_radius_2d,
             eps_2d=eps_2d,
             antialias=antialias,
+            backgrounds=backgrounds,
         )
 
     def render_images_and_depths(
@@ -1791,6 +1797,7 @@ def render_images_and_depths(
         min_radius_2d: float = 0.0,
         eps_2d: float = 0.3,
         antialias: bool = False,
+        backgrounds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Render ``C`` multi-channel images (see :attr:`num_channels`) with depth as the last channel from this :class:`GaussianSplat3d` from ``C`` camera views.
@@ -1861,6 +1868,7 @@ def render_images_and_depths(
             min_radius_2d=min_radius_2d,
             eps_2d=eps_2d,
             antialias=antialias,
+            backgrounds=backgrounds,
         )
 
     def render_num_contributing_gaussians(