Skip to content

Commit 5988e39

Browse files
matthewdcongmatthewdcong
committed
Add Gaussian isocontour tile intersection
Signed-off-by: Matthew Cong <mcong@nvidia.com>
1 parent a4ee8dd commit 5988e39

1 file changed

Lines changed: 95 additions & 34 deletions

File tree

‎fvdb_reality_capture/radiance_fields/gaussian_splatting.py‎

Lines changed: 95 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1667,6 +1667,8 @@ def _intersect_tiles(
16671667
means2d: torch.Tensor,
16681668
radii: torch.Tensor,
16691669
depths: torch.Tensor,
1670+
conics: torch.Tensor,
1671+
opacities: torch.Tensor,
16701672
C: int,
16711673
tile_size: int,
16721674
W: int,
@@ -1686,6 +1688,8 @@ def _intersect_tiles(
16861688
tile_size,
16871689
num_tiles_h,
16881690
num_tiles_w,
1691+
conics=conics,
1692+
opacities=opacities,
16891693
)
16901694
return tile_offsets, tile_gaussian_ids, num_tiles_h, num_tiles_w
16911695

@@ -1695,6 +1699,8 @@ def _intersect_tiles_sparse(
16951699
means2d: torch.Tensor,
16961700
radii: torch.Tensor,
16971701
depths: torch.Tensor,
1702+
conics: torch.Tensor,
1703+
opacities: torch.Tensor,
16981704
C: int,
16991705
tile_size: int,
17001706
W: int,
@@ -1706,13 +1712,17 @@ def _intersect_tiles_sparse(
17061712
"""
17071713
num_tiles_h = math.ceil(H / tile_size)
17081714
num_tiles_w = math.ceil(W / tile_size)
1709-
active_tiles, active_tile_mask, tile_pixel_mask, tile_pixel_cumsum, pixel_map = (
1710-
_C.build_sparse_gaussian_tile_layout(
1711-
tile_size,
1712-
num_tiles_w,
1713-
num_tiles_h,
1714-
pixels_jt._impl,
1715-
)
1715+
(
1716+
active_tiles,
1717+
active_tile_mask,
1718+
tile_pixel_mask,
1719+
tile_pixel_cumsum,
1720+
pixel_map,
1721+
) = _C.build_sparse_gaussian_tile_layout(
1722+
tile_size,
1723+
num_tiles_w,
1724+
num_tiles_h,
1725+
pixels_jt._impl,
17161726
)
17171727
tile_offsets, tile_gaussian_ids = _C.intersect_gaussian_tiles_sparse(
17181728
means2d,
@@ -1724,6 +1734,8 @@ def _intersect_tiles_sparse(
17241734
tile_size,
17251735
num_tiles_h,
17261736
num_tiles_w,
1737+
conics=conics,
1738+
opacities=opacities,
17271739
)
17281740
return tile_offsets, tile_gaussian_ids, active_tiles, tile_pixel_mask, tile_pixel_cumsum, pixel_map
17291741

@@ -1961,9 +1973,14 @@ def _sparse_render_impl(
19611973
opacities = self._make_opacities(C, compensations, antialias)
19621974
features = self._make_render_features(w2c, radii, depths, sh_degree_to_use, include_colors, include_depth)
19631975

1964-
tile_offsets, tile_gaussian_ids, active_tiles, tile_pixel_mask, tile_pixel_cumsum, pixel_map = (
1965-
self._intersect_tiles_sparse(render_pixels, means2d, radii, depths, C, tile_size, W, H)
1966-
)
1976+
(
1977+
tile_offsets,
1978+
tile_gaussian_ids,
1979+
active_tiles,
1980+
tile_pixel_mask,
1981+
tile_pixel_cumsum,
1982+
pixel_map,
1983+
) = self._intersect_tiles_sparse(render_pixels, means2d, radii, depths, conics, opacities, C, tile_size, W, H)
19671984

19681985
rendered_jdata, alphas_jdata = self._rasterize_screen_space_sparse(
19691986
render_pixels,
@@ -2503,6 +2520,8 @@ def render_from_projected_gaussians(
25032520
tile_size,
25042521
num_tiles_h,
25052522
num_tiles_w,
2523+
conics=pg.inv_covar_2d,
2524+
opacities=pg.opacities,
25062525
)
25072526
features, alphas = cast(
25082527
tuple[torch.Tensor, torch.Tensor],
@@ -2528,6 +2547,8 @@ def render_from_projected_gaussians(
25282547
pg.means2d,
25292548
pg.radii,
25302549
pg.depths,
2550+
pg.inv_covar_2d,
2551+
pg.opacities,
25312552
C,
25322553
tile_size,
25332554
W,
@@ -2655,6 +2676,8 @@ def render_depths(
26552676
means2d,
26562677
radii,
26572678
depths,
2679+
conics,
2680+
opacities,
26582681
C,
26592682
tile_size,
26602683
image_width,
@@ -2903,6 +2926,8 @@ def render_images(
29032926
means2d,
29042927
radii,
29052928
depths,
2929+
conics,
2930+
opacities,
29062931
C,
29072932
tile_size,
29082933
image_width,
@@ -3041,6 +3066,8 @@ def render_images_from_world(
30413066
means2d,
30423067
radii,
30433068
depths,
3069+
conics,
3070+
opacities,
30443071
C,
30453072
tile_size,
30463073
image_width,
@@ -3115,6 +3142,8 @@ def render_depths_from_world(
31153142
means2d,
31163143
radii,
31173144
depths,
3145+
conics,
3146+
opacities,
31183147
C,
31193148
tile_size,
31203149
image_width,
@@ -3509,6 +3538,8 @@ def render_images_and_depths(
35093538
means2d,
35103539
radii,
35113540
depths,
3541+
conics,
3542+
opacities,
35123543
C,
35133544
tile_size,
35143545
image_width,
@@ -3585,6 +3616,8 @@ def render_images_and_depths_from_world(
35853616
means2d,
35863617
radii,
35873618
depths,
3619+
conics,
3620+
opacities,
35883621
C,
35893622
tile_size,
35903623
image_width,
@@ -3711,6 +3744,8 @@ def render_num_contributing_gaussians(
37113744
means2d,
37123745
radii,
37133746
depths,
3747+
conics,
3748+
opacities,
37143749
C,
37153750
tile_size,
37163751
image_width,
@@ -3857,17 +3892,24 @@ def sparse_render_num_contributing_gaussians(
38573892
)
38583893
C = world_to_camera_matrices.size(0)
38593894
opacities = self._make_opacities(C, compensations, antialias)
3860-
tile_offsets, tile_gaussian_ids, active_tiles, tile_pixel_mask, tile_pixel_cumsum, pixel_map = (
3861-
self._intersect_tiles_sparse(
3862-
unique_pixels_jt,
3863-
means2d,
3864-
radii,
3865-
depths,
3866-
C,
3867-
tile_size,
3868-
image_width,
3869-
image_height,
3870-
)
3895+
(
3896+
tile_offsets,
3897+
tile_gaussian_ids,
3898+
active_tiles,
3899+
tile_pixel_mask,
3900+
tile_pixel_cumsum,
3901+
pixel_map,
3902+
) = self._intersect_tiles_sparse(
3903+
unique_pixels_jt,
3904+
means2d,
3905+
radii,
3906+
depths,
3907+
conics,
3908+
opacities,
3909+
C,
3910+
tile_size,
3911+
image_width,
3912+
image_height,
38713913
)
38723914
result_ncg, result_alphas = _C.sparse_rasterize_num_contributing_gaussians(
38733915
means2d,
@@ -3983,6 +4025,8 @@ def render_contributing_gaussian_ids(
39834025
means2d,
39844026
radii,
39854027
depths,
4028+
conics,
4029+
opacities,
39864030
C,
39874031
tile_size,
39884032
image_width,
@@ -4144,17 +4188,24 @@ def sparse_render_contributing_gaussian_ids(
41444188
)
41454189
C = world_to_camera_matrices.size(0)
41464190
opacities = self._make_opacities(C, compensations, antialias)
4147-
tile_offsets, tile_gaussian_ids, active_tiles, tile_pixel_mask, tile_pixel_cumsum, pixel_map = (
4148-
self._intersect_tiles_sparse(
4149-
unique_pixels_jt,
4150-
means2d,
4151-
radii,
4152-
depths,
4153-
C,
4154-
tile_size,
4155-
image_width,
4156-
image_height,
4157-
)
4191+
(
4192+
tile_offsets,
4193+
tile_gaussian_ids,
4194+
active_tiles,
4195+
tile_pixel_mask,
4196+
tile_pixel_cumsum,
4197+
pixel_map,
4198+
) = self._intersect_tiles_sparse(
4199+
unique_pixels_jt,
4200+
means2d,
4201+
radii,
4202+
depths,
4203+
conics,
4204+
opacities,
4205+
C,
4206+
tile_size,
4207+
image_width,
4208+
image_height,
41584209
)
41594210
ncg_jt = None
41604211
if top_k_contributors <= 0:
@@ -4669,6 +4720,7 @@ def gaussian_render_jagged(
46694720
opacities_batched = opacities.jdata[gaussian_ids]
46704721
if antialias:
46714722
opacities_batched = opacities_batched * compensations
4723+
opacities_batched = opacities_batched.contiguous()
46724724

46734725
debug_info: dict[str, torch.Tensor] = {}
46744726
if return_debug_info:
@@ -4723,7 +4775,16 @@ def gaussian_render_jagged(
47234775
num_tiles_h = math.ceil(image_height / tile_size)
47244776
num_tiles_w = math.ceil(image_width / tile_size)
47254777
tile_offsets, tile_gaussian_ids_t = _C.intersect_gaussian_tiles(
4726-
means2d, radii, depths, ccz, tile_size, num_tiles_h, num_tiles_w, camera_ids
4778+
means2d,
4779+
radii,
4780+
depths,
4781+
ccz,
4782+
tile_size,
4783+
num_tiles_h,
4784+
num_tiles_w,
4785+
camera_ids=camera_ids,
4786+
conics=conics,
4787+
opacities=opacities_batched,
47274788
)
47284789
if return_debug_info:
47294790
debug_info["tile_offsets"] = tile_offsets
@@ -4734,7 +4795,7 @@ def gaussian_render_jagged(
47344795
means2d,
47354796
conics,
47364797
render_quantities,
4737-
opacities_batched.contiguous(),
4798+
opacities_batched,
47384799
image_width,
47394800
image_height,
47404801
0, # image_origin_w

0 commit comments

Comments
 (0)