From a03fd266c6322510caab232f8b323e23c82609cd Mon Sep 17 00:00:00 2001 From: Marcin Ochman Date: Mon, 3 Jun 2024 14:56:26 +0200 Subject: [PATCH 1/3] Viewer support for ds-splat --- internal/renderers/.gitignore | 3 +- internal/renderers/ds_splat_renderer.py | 198 ++++++++++++++++++++++++ viewer.py | 8 + 3 files changed, 208 insertions(+), 1 deletion(-) create mode 100644 internal/renderers/ds_splat_renderer.py diff --git a/internal/renderers/.gitignore b/internal/renderers/.gitignore index 9662994e..35ea2a8b 100644 --- a/internal/renderers/.gitignore +++ b/internal/renderers/.gitignore @@ -12,4 +12,5 @@ !pypreprocess_gsplat_renderer.py !swag_renderer.py !mip_splatting_gsplat_renderer.py -!gsplat_hit_pixel_count_renderer.py \ No newline at end of file +!gsplat_hit_pixel_count_renderer.py +!ds_splat_renderer.py diff --git a/internal/renderers/ds_splat_renderer.py b/internal/renderers/ds_splat_renderer.py new file mode 100644 index 00000000..a1447dbc --- /dev/null +++ b/internal/renderers/ds_splat_renderer.py @@ -0,0 +1,198 @@ +import math +import torch +from .renderer import Renderer, Camera, GaussianModel +from ds_splat import ( + GaussianRasterizationSettings, + GaussianRasterizer, +) +from typing import Optional +from internal.utils.sh_utils import eval_sh + + +class OpenRenderer(Renderer): + def __init__( + self, precompute_cov_3d: bool = False, precompute_colors: bool = False + ): + super().__init__() + + self.precompute_cov_3d = precompute_cov_3d + self.precompute_colors = precompute_colors + + def forward( + self, + viewpoint_camera: Camera, + pc: GaussianModel, + bg_color: torch.Tensor, + scaling_modifier=1.0, + override_color=None, + ): + screenspace_points = ( + torch.zeros_like( + pc.get_xyz, + dtype=pc.get_xyz.dtype, + requires_grad=True, + device=bg_color.device, + ) + + 0 + ) + + try: + screenspace_points.retain_grad() + except RuntimeError: + pass + + tanfovx = math.tan(viewpoint_camera.fov_x * 0.5) + tanfovy = math.tan(viewpoint_camera.fov_y * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.height), + image_width=int(viewpoint_camera.width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_to_camera, + projmatrix=viewpoint_camera.full_projection, + sh_degree=pc.active_sh_degree, + max_sh_degree=3, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=False, + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means_3d = pc.get_xyz + means_2d = screenspace_points + opacity = pc.get_opacity + + scales, rotations, cov_3d_precomp, shs, colors_precomp = ( + self.get_cov_3d_and_colors( + viewpoint_camera, pc, scaling_modifier, override_color + ) + ) + + rendered_image, radii = rasterizer( + means_3d=means_3d, + means_2d=means_2d, + shs=shs, + colors_precomp=colors_precomp, + opacities=opacity, + scales=scales, + rotations=rotations, + cov_3d_precomp=cov_3d_precomp, + ) + + grad_scale = 0.5 * max( + raster_settings.image_height, raster_settings.image_width + ) + + return { + "render": rendered_image.permute(2, 0, 1), + "viewspace_points": screenspace_points, + "viewspace_points_grad_scale": grad_scale, + "visibility_filter": radii > 0, + "radii": radii, + } + + def get_cov_3d_and_colors( + self, + viewpoint_camera: Camera, + pc: GaussianModel, + scaling_modifier=1.0, + override_color=None, + ): + scales = None + rotations = None + cov_3d_precomp = None + + if self.precompute_cov_3d is True: + cov_3d_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = None + colors_precomp = None + if override_color is None: + if self.precompute_colors is True: + shs_view = pc.get_features.transpose(1, 2).view( + -1, 3, (pc.max_sh_degree + 1) ** 2 + ) + dir_pp = pc.get_xyz - viewpoint_camera.camera_center.repeat( + pc.get_features.shape[0], 1 + ) + dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True) + sh2_to_rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized) + + colors_precomp = torch.clamp_min(sh2_to_rgb + 0.5, 0.0) + else: + shs = pc.get_features + else: + colors_precomp = override_color + + return scales, rotations, cov_3d_precomp, shs, colors_precomp + + @staticmethod + def render( + means_3d: torch.Tensor, + opacity: torch.Tensor, + scales: Optional[torch.Tensor], + rotations: Optional[torch.Tensor], + features: Optional[torch.Tensor], + active_sh_degree: int, + viewpoint_camera, + bg_color: torch.Tensor, + scaling_modifier=1.0, + colors_precomp: Optional[torch.Tensor] = None, + cov_3d_precomp: Optional[torch.Tensor] = None, + ): + screenspace_points = torch.zeros_like( + means_3d, + dtype=means_3d.dtype, + requires_grad=True, + device=means_3d.device, + ) + + try: + screenspace_points.retain_grad() + except RuntimeError: + pass + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.height), + image_width=int(viewpoint_camera.width), + tanfovx=math.tan(viewpoint_camera.fov_x * 0.5), + tanfovy=math.tan(viewpoint_camera.fov_y * 0.5), + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_to_camera, + projmatrix=viewpoint_camera.full_projection, + sh_degree=active_sh_degree, + max_sh_degree=3, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=False, + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + means_2d = screenspace_points + + rendered_image, radii = rasterizer( + means_3d=means_3d, + means_2d=means_2d, + shs=features, + colors_precomp=colors_precomp, + opacities=opacity, + scales=scales, + rotations=rotations, + cov_3d_precomp=cov_3d_precomp, + ) + + return { + "render": rendered_image.permute(2, 0, 1), + "depth": None, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } diff --git a/viewer.py b/viewer.py index 9c2443d3..944b5dbe 100644 --- a/viewer.py +++ b/viewer.py @@ -42,6 +42,7 @@ def __init__( no_edit_panel: bool = False, no_render_panel: bool = False, gsplat: bool = False, + dssplat: bool = False, ): self.device = torch.device("cuda") @@ -60,6 +61,7 @@ def __init__( self.default_camera_look_at = default_camera_look_at self.use_gsplat = gsplat + self.use_dssplat = dssplat load_from = self._search_load_file(model_paths[0]) @@ -311,6 +313,10 @@ def _load_model_from_file(self, load_from: str): from internal.renderers.gsplat_renderer import GSPlatRenderer print("Use GSPlat renderer for ply file") renderer = GSPlatRenderer() + elif self.use_dssplat is True: + from internal.renderers.ds_splat_renderer import OpenRenderer + print("Use ds_splat renderer for ply file") + renderer = OpenRenderer() else: raise ValueError("unsupported file {}".format(load_from)) @@ -606,6 +612,8 @@ def _handle_client_disconnect(self, client: viser.ClientHandle): parser.add_argument("--no_render_panel", action="store_true", default=False) parser.add_argument("--gsplat", action="store_true", default=False, help="Use GSPlat renderer for ply file") + parser.add_argument("--dssplat", action="store_true", default=False, + help="Use ds_splat renderer for ply file") parser.add_argument("--float32_matmul_precision", "--fp", type=str, default=None) args = parser.parse_args() From 7e2f1ce0ee3f749267cf5b009096345976b592a3 Mon Sep 17 00:00:00 2001 From: Marcin Ochman Date: Mon, 3 Jun 2024 15:34:14 +0200 Subject: [PATCH 2/3] Added ds_splat config --- configs/.gitignore | 3 ++- configs/ds_splat.yaml | 2 ++ 2 files changed, 4 insertions(+), 1 deletion(-) create mode 100644 configs/ds_splat.yaml diff --git a/configs/.gitignore b/configs/.gitignore index bcc50a9c..d4ec0911 100644 --- a/configs/.gitignore +++ b/configs/.gitignore @@ -23,4 +23,5 @@ !light_gaussian/train_densify_prune-gsplat.yaml !light_gaussian/train_densify_prune-gsplat-experiment.yaml !gsplat-absgrad.yaml -!gsplat-absgrad-experiment.yaml \ No newline at end of file +!gsplat-absgrad-experiment.yaml +!ds_splat.yaml diff --git a/configs/ds_splat.yaml b/configs/ds_splat.yaml new file mode 100644 index 00000000..d3f503fa --- /dev/null +++ b/configs/ds_splat.yaml @@ -0,0 +1,2 @@ +model: + renderer: internal.renderers.ds_splat_renderer.OpenRenderer From c03d809b77aaf2a8155d55dfcfae61aad56bedae Mon Sep 17 00:00:00 2001 From: Marcin Ochman Date: Mon, 1 Jul 2024 14:20:58 +0200 Subject: [PATCH 3/3] Changed to DsRenderer --- README.md | 5 +++++ configs/ds_splat.yaml | 2 +- internal/renderers/ds_splat_renderer.py | 2 +- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index ad0db82b..73804d59 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,11 @@ pip install -r requirements.txt ```bash pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch ``` +* If you want to use ds-splat: + + ```bash + pip install ds-splat==0.0.1 + ``` * If you want to use nerfstudio-project/gsplat * Vanilla version diff --git a/configs/ds_splat.yaml b/configs/ds_splat.yaml index d3f503fa..7fb82d81 100644 --- a/configs/ds_splat.yaml +++ b/configs/ds_splat.yaml @@ -1,2 +1,2 @@ model: - renderer: internal.renderers.ds_splat_renderer.OpenRenderer + renderer: internal.renderers.ds_splat_renderer.DsRenderer diff --git a/internal/renderers/ds_splat_renderer.py b/internal/renderers/ds_splat_renderer.py index a1447dbc..44949a3d 100644 --- a/internal/renderers/ds_splat_renderer.py +++ b/internal/renderers/ds_splat_renderer.py @@ -9,7 +9,7 @@ from internal.utils.sh_utils import eval_sh -class OpenRenderer(Renderer): +class DsRenderer(Renderer): def __init__( self, precompute_cov_3d: bool = False, precompute_colors: bool = False ):