Ds splat support

README.md

@@ -75,6 +75,7 @@ pip install -r requirements.txt
 ```
 
 ### 1.5. Install optional packages
+
 * <a href="https://ffmpeg.org/">ffmpeg</a> is required if you want to render video: `sudo apt install -y ffmpeg`
 * If you want to use <a href="https://github.com/nerfstudio-project/gsplat">nerfstudio-project/gsplat</a>
 
@@ -83,7 +84,12 @@ pip install -r requirements.txt
     ```
 
   This command will install my modified version, which is required by LightGaussian and Mip-Splatting. If you do not need them, you can also install vanilla gsplat <a href="https://github.com/nerfstudio-project/gsplat/tree/v0.1.12">v0.1.12</a>.
-
+
+* If you want to use ds-splat:
+
+    ```bash
+    pip install ds-splat==0.0.1
+    ```
 * If you need <a href="#210-segment-any-3d-gaussians">SegAnyGaussian</a>
   * gsplat (see command above)
   * `pip install hdbscan scikit-learn==1.3.2 git+https://github.com/facebookresearch/segment-anything.git`

configs/.gitignore

@@ -25,10 +25,11 @@
 !light_gaussian/train_densify_prune-gsplat-experiment.yaml
 !gsplat-absgrad.yaml
 !gsplat-absgrad-experiment.yaml
+!ds_splat.yaml
 !vanilla_2dgs.yaml
 !segany_splatting.yaml
 !appearance_embedding_renderer/view_dependent.yaml
 !appearance_embedding_renderer/view_independent.yaml
 !random_background.yaml
 !mcmc.yaml
-!gsplat-mcmc.yaml
+!gsplat-mcmc.yaml

configs/ds_splat.yaml

@@ -0,0 +1,2 @@
+model:
+  renderer: internal.renderers.ds_splat_renderer.DsRenderer

internal/renderers/.gitignore

@@ -13,6 +13,7 @@
 !swag_renderer.py
 !mip_splatting_gsplat_renderer.py
 !gsplat_hit_pixel_count_renderer.py
+!ds_splat_renderer.py
 !vanilla_2dgs_renderer.py
 !seganygs_renderer.py
-!gsplat_appearance_embedding_renderer.py
+!gsplat_appearance_embedding_renderer.py

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 DsRenderer(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,
+        }

viewer.py

@@ -43,6 +43,7 @@ def __init__(
             no_edit_panel: bool = False,
             no_render_panel: bool = False,
             gsplat: bool = False,
+            dssplat: bool = False,
             seganygs: str = None,
             vanilla_seganygs: bool = False,
             vanilla_mip: bool = False,
@@ -65,6 +66,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])
 
@@ -424,6 +426,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))
 
@@ -745,6 +751,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("--seganygs", type=str, default=None,
                         help="Path to a SegAnyGaussian model output directory or checkpoint file")
     parser.add_argument("--vanilla_seganygs", action="store_true", default=False)