reduced-3dgs 1.8.14__tar.gz → 1.8.16__tar.gz

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reduced_3dgs
-Version: 1.8.14
+Version: 1.8.16
 Summary: Refactored code for the paper "Reducing the Memory Footprint of 3D Gaussian Splatting"
 Home-page: https://github.com/yindaheng98/reduced-3dgs
 Author: yindaheng98

reduced_3dgs-1.8.16/reduced_3dgs/pruning/__init__.py

@@ -0,0 +1,2 @@
+from .trainer import BasePruner, BasePruningTrainer
+from .combinations import BasePrunerInDensifyTrainer, PruningTrainer, PrunerInDensifyTrainer

reduced_3dgs-1.8.16/reduced_3dgs/pruning/combinations.py

@@ -0,0 +1,51 @@
+
+from typing import List
+from gaussian_splatting import Camera, GaussianModel
+from gaussian_splatting.dataset import TrainableCameraDataset
+from gaussian_splatting.trainer import DepthTrainerWrapper, NoopDensifier, DensificationTrainerWrapper
+from .trainer import BasePruner, BasePruningTrainer
+
+
+def BasePrunerInDensifyTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        prune_from_iter=1000,
+        prune_until_iter=15000,
+        prune_interval: int = 100,
+        box_size=1.,
+        lambda_mercy=1.,
+        mercy_minimum=3,
+        mercy_type='redundancy_opacity',
+        *args, **kwargs):
+    return DensificationTrainerWrapper(
+        lambda model, scene_extent: BasePruner(
+            NoopDensifier(model),
+            model, dataset,
+            prune_from_iter=prune_from_iter,
+            prune_until_iter=prune_until_iter,
+            prune_interval=prune_interval,
+            box_size=box_size,
+            lambda_mercy=lambda_mercy,
+            mercy_minimum=mercy_minimum,
+            mercy_type=mercy_type,
+        ),
+        model,
+        scene_extent,
+        *args, **kwargs
+    )
+
+
+# Depth trainer
+
+
+def DepthPruningTrainer(model: GaussianModel, scene_extent: float, dataset: TrainableCameraDataset, *args, **kwargs):
+    return DepthTrainerWrapper(BasePruningTrainer, model, scene_extent, *args, dataset=dataset, **kwargs)
+
+
+def DepthPrunerInDensifyTrainer(model: GaussianModel, scene_extent: float, dataset: TrainableCameraDataset, *args, **kwargs):
+    return DepthTrainerWrapper(BasePrunerInDensifyTrainer, model, scene_extent, *args, dataset=dataset, **kwargs)
+
+
+PruningTrainer = DepthPruningTrainer
+PrunerInDensifyTrainer = DepthPrunerInDensifyTrainer

reduced_3dgs-1.8.16/reduced_3dgs/pruning/importance/__init__.py

@@ -0,0 +1 @@
+from .trainer import ImportancePruningTrainerWrapper, BaseImportancePruningTrainer, ImportancePruningTrainer

reduced_3dgs-1.8.16/reduced_3dgs/pruning/importance/trainer.py

@@ -0,0 +1,141 @@
+import math
+import torch
+
+from gaussian_splatting import Camera, GaussianModel
+from gaussian_splatting.trainer import AbstractTrainer, TrainerWrapper, BaseTrainer, Trainer
+from gaussian_splatting.dataset import CameraDataset
+from .diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
+
+
+def count_render(self: GaussianModel, viewpoint_camera: Camera):
+    """
+    Render the scene.
+
+    Background tensor (bg_color) must be on GPU!
+    """
+    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
+    screenspace_points = torch.zeros_like(self.get_xyz, dtype=self.get_xyz.dtype, requires_grad=True, device=self._xyz.device) + 0
+    try:
+        screenspace_points.retain_grad()
+    except:
+        pass
+
+    # Set up rasterization configuration
+    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
+    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
+
+    raster_settings = GaussianRasterizationSettings(
+        image_height=int(viewpoint_camera.image_height),
+        image_width=int(viewpoint_camera.image_width),
+        tanfovx=tanfovx,
+        tanfovy=tanfovy,
+        bg=viewpoint_camera.bg_color.to(self._xyz.device),
+        scale_modifier=self.scale_modifier,
+        viewmatrix=viewpoint_camera.world_view_transform,
+        projmatrix=viewpoint_camera.full_proj_transform,
+        sh_degree=self.active_sh_degree,
+        campos=viewpoint_camera.camera_center,
+        prefiltered=False,
+        debug=self.debug,
+        f_count=True,
+    )
+
+    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
+    means3D = self.get_xyz
+    means2D = screenspace_points
+    opacity = self.get_opacity
+
+    scales = self.get_scaling
+    rotations = self.get_rotation
+
+    shs = self.get_features
+
+    # Rasterize visible Gaussians to image, obtain their radii (on screen).
+    gaussians_count, opacity_important_score, T_alpha_important_score, rendered_image, radii = rasterizer(
+        means3D=means3D,
+        means2D=means2D,
+        shs=shs,
+        colors_precomp=None,
+        opacities=opacity,
+        scales=scales,
+        rotations=rotations,
+        cov3D_precomp=None)
+
+    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
+    # They will be excluded from value updates used in the splitting criteria.
+    return {
+        "render": rendered_image,
+        "viewspace_points": screenspace_points,
+        "visibility_filter": radii > 0,
+        "radii": radii,
+        "gaussians_count": gaussians_count,
+        "opacity_important_score": opacity_important_score,
+        "T_alpha_important_score": T_alpha_important_score
+    }
+
+
+class ImportancePruner(TrainerWrapper):
+    def __init__(
+            self, base_trainer: AbstractTrainer,
+            dataset: CameraDataset,
+            importance_prune_at_steps=[15000],
+    ):
+        super().__init__(base_trainer)
+        self.dataset = dataset
+        self.importance_prune_at_steps = importance_prune_at_steps
+
+    def optim_step(self):
+        ret = super().optim_step()
+        if self.curr_step in self.importance_prune_at_steps:
+            gaussian_count = torch.zeros(self.model.get_xyz.shape[0], device=self.model.get_xyz.device, dtype=torch.int)
+            opacity_important_score = torch.zeros(self.model.get_xyz.shape[0], device=self.model.get_xyz.device, dtype=torch.float)
+            T_alpha_important_score = torch.zeros(self.model.get_xyz.shape[0], device=self.model.get_xyz.device, dtype=torch.float)
+            for camera in self.dataset:
+                out = count_render(self.model, camera)
+                gaussian_count += out["gaussians_count"]
+                opacity_important_score += out["opacity_important_score"]
+                T_alpha_important_score += out["T_alpha_important_score"]
+            pass
+        return ret
+
+
+def ImportancePruningTrainerWrapper(
+    base_trainer_constructor,
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        importance_prune_at_steps=[15000],
+        *args, **kwargs):
+    return ImportancePruner(
+        base_trainer_constructor(model, scene_extent, dataset, *args, **kwargs),
+        dataset,
+        importance_prune_at_steps=importance_prune_at_steps,
+    )
+
+
+def BaseImportancePruningTrainer(
+    model: GaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        importance_prune_at_steps=[15000],
+        *args, **kwargs):
+    return ImportancePruningTrainerWrapper(
+        lambda model, scene_extent, dataset, *args, **kwargs: BaseTrainer(model, scene_extent, *args, **kwargs),
+        model, scene_extent, dataset,
+        importance_prune_at_steps=importance_prune_at_steps,
+        *args, **kwargs,
+    )
+
+
+def ImportancePruningTrainer(
+    model: GaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        importance_prune_at_steps=[15000],
+        *args, **kwargs):
+    return ImportancePruningTrainerWrapper(
+        lambda model, scene_extent, dataset, *args, **kwargs: Trainer(model, scene_extent, *args, **kwargs),
+        model, scene_extent, dataset,
+        importance_prune_at_steps=importance_prune_at_steps,
+        *args, **kwargs,
+    )

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/reduced_3dgs/pruning/trainer.py

@@ -1,7 +1,7 @@
 from typing import List
 import torch
 from gaussian_splatting import GaussianModel, Camera
-from gaussian_splatting.trainer import AbstractDensifier, Densifier, DensificationInstruct, DensificationTrainer
+from gaussian_splatting.trainer import AbstractDensifier, DensifierWrapper, DensificationTrainer, NoopDensifier
 from reduced_3dgs.diff_gaussian_rasterization._C import sphere_ellipsoid_intersection, allocate_minimum_redundancy_value, find_minimum_projected_pixel_size
 from reduced_3dgs.simple_knn._C import distIndex2
 
@@ -79,16 +79,18 @@ def mercy_gaussians(
     return mask
 
 
-class BasePruner(AbstractDensifier):
+class BasePruner(DensifierWrapper):
     def __init__(
-            self, model: GaussianModel, dataset: List[Camera],
+            self, base_densifier: AbstractDensifier,
+            model: GaussianModel, dataset: List[Camera],
             prune_from_iter=1000,
             prune_until_iter=15000,
-            prune_interval: int = 500,
+            prune_interval: int = 100,
             box_size=1.,
             lambda_mercy=1.,
             mercy_minimum=3,
             mercy_type='redundancy_opacity'):
+        super().__init__(base_densifier)
         self._model = model
         self.dataset = dataset
         self.prune_from_iter = prune_from_iter
@@ -103,10 +105,12 @@ class BasePruner(AbstractDensifier):
     def model(self) -> GaussianModel:
         return self._model
 
-    def densify_and_prune(self, loss, out, camera, step: int) -> DensificationInstruct:
-        if self.prune_from_iter < step < self.prune_until_iter and step % self.prune_interval == 0:
-            return DensificationInstruct(remove_mask=mercy_gaussians(self.model, self.dataset, self.box_size, self.lambda_mercy, self.mercy_minimum, self.mercy_type))
-        return DensificationInstruct()
+    def densify_and_prune(self, loss, out, camera, step: int):
+        ret = super().densify_and_prune(loss, out, camera, step)
+        if self.prune_from_iter <= step <= self.prune_until_iter and step % self.prune_interval == 0:
+            remove_mask = mercy_gaussians(self.model, self.dataset, self.box_size, self.lambda_mercy, self.mercy_minimum, self.mercy_type)
+            ret = ret._replace(remove_mask=remove_mask if ret.remove_mask is None else torch.logical_or(remove_mask, ret.remove_mask))
+        return ret
 
 
 def BasePruningTrainer(
@@ -115,7 +119,7 @@ def BasePruningTrainer(
         dataset: List[Camera],
         prune_from_iter=1000,
         prune_until_iter=15000,
-        prune_interval: int = 500,
+        prune_interval: int = 100,
         box_size=1.,
         lambda_mercy=1.,
         mercy_minimum=3,
@@ -124,75 +128,14 @@ def BasePruningTrainer(
     return DensificationTrainer(
         model, scene_extent,
         BasePruner(
+            NoopDensifier(),
             model, dataset,
-            prune_from_iter, prune_until_iter, prune_interval,
-            box_size, lambda_mercy, mercy_minimum, mercy_type
-        ), *args, **kwargs
-    )
-
-
-class PrunerInDensify(Densifier):
-    def __init__(
-            self, model: GaussianModel, scene_extent, dataset: List[Camera],
-            box_size=1.,
-            lambda_mercy=1.,
-            mercy_minimum=3,
-            mercy_type='redundancy_opacity',
-            *args, **kwargs):
-        super().__init__(model, scene_extent, *args, **kwargs)
-        self.dataset = dataset
-        self.box_size = box_size
-        self.lambda_mercy = lambda_mercy
-        self.mercy_minimum = mercy_minimum
-        self.mercy_type = mercy_type
-
-    def prune(self) -> torch.Tensor:
-        _splatted_num_accum, _ = calculate_redundancy_metric(self.model, self.dataset, pixel_scale=self.box_size)
-        mask = mercy_points(self.model, _splatted_num_accum.squeeze(), self.lambda_mercy, self.mercy_minimum, self.mercy_type)
-        return torch.logical_or(mask, super().prune())
-
-
-def BasePrunerInDensifyTrainer(
-        model: GaussianModel,
-        scene_extent: float,
-
-        dataset: List[Camera],
-        box_size=1.,
-        lambda_mercy=1.,
-        mercy_minimum=3,
-        mercy_type='redundancy_opacity',
-
-        densify_from_iter=500,
-        densify_until_iter=15000,
-        densify_interval=100,
-        densify_grad_threshold=0.0002,
-        densify_opacity_threshold=0.005,
-        densify_percent_dense=0.01,
-        densify_percent_too_big=0.8,
-
-        prune_from_iter=1000,
-        prune_until_iter=15000,
-        prune_interval=500,
-        prune_screensize_threshold=20,
-        prune_percent_too_big=1,
-
-        *args, **kwargs):
-    return DensificationTrainer(
-        model, scene_extent,
-        PrunerInDensify(
-            model, scene_extent, dataset,
-            box_size, lambda_mercy, mercy_minimum, mercy_type,
-            densify_from_iter=densify_from_iter,
-            densify_until_iter=densify_until_iter,
-            densify_interval=densify_interval,
-            densify_grad_threshold=densify_grad_threshold,
-            densify_opacity_threshold=densify_opacity_threshold,
-            densify_percent_dense=densify_percent_dense,
-            densify_percent_too_big=densify_percent_too_big,
             prune_from_iter=prune_from_iter,
             prune_until_iter=prune_until_iter,
             prune_interval=prune_interval,
-            prune_screensize_threshold=prune_screensize_threshold,
-            prune_percent_too_big=prune_percent_too_big
+            box_size=box_size,
+            lambda_mercy=lambda_mercy,
+            mercy_minimum=mercy_minimum,
+            mercy_type=mercy_type,
         ), *args, **kwargs
     )

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/reduced_3dgs/quantization/abc.py

@@ -42,7 +42,7 @@ class QuantizeTrainerWrapper(TrainerWrapper, metaclass=abc.ABCMeta):
 
     @property
     def model(self) -> GaussianModel:
-        if self.quantize_from_iter <= self.curr_step < self.quantize_until_iter and self.curr_step % self.quantize_interval == 0:
+        if self.quantize_from_iter <= self.curr_step <= self.quantize_until_iter and self.curr_step % self.quantize_interval == 0:
             with torch.no_grad():
                 ids_dict, codebook_dict = self.quantizer.quantize(self.base_trainer.model, update_codebook=True)
                 return self.quantizer.dequantize(self.base_trainer.model, ids_dict, codebook_dict)

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/reduced_3dgs/train.py

@@ -14,6 +14,7 @@ from gaussian_splatting.trainer.extensions import ScaleRegularizeTrainerWrapper
 from reduced_3dgs.quantization import AbstractQuantizer, VectorQuantizeTrainerWrapper
 from reduced_3dgs.shculling import VariableSHGaussianModel, SHCullingTrainer
 from reduced_3dgs.pruning import PruningTrainer
+from reduced_3dgs.pruning.importance import ImportancePruningTrainerWrapper
 from reduced_3dgs.combinations import OpacityResetPrunerInDensifyTrainer, SHCullingDensifyTrainer, SHCullingPruneTrainer, SHCullingPruningDensifyTrainer
 from reduced_3dgs.combinations import CameraTrainableVariableSHGaussianModel, CameraSHCullingTrainer, CameraPruningTrainer
 from reduced_3dgs.combinations import CameraOpacityResetPrunerInDensifyTrainer, CameraSHCullingDensifyTrainer, CameraSHCullingPruneTrainer, CameraSHCullingPruningDensifyTrainer
@@ -27,6 +28,7 @@ basemodes = {
     "prune-shculling": SHCullingPruneTrainer,
     "densify-prune-shculling": SHCullingPruningDensifyTrainer,
 }
+basemodes = {k: lambda *args, **kwargs: ImportancePruningTrainerWrapper(v, *args, **kwargs) for k, v in basemodes.items()}
 cameramodes = {
     "camera-shculling": CameraSHCullingTrainer,
     "camera-pruning": CameraPruningTrainer,
@@ -35,6 +37,7 @@ cameramodes = {
     "camera-prune-shculling": CameraSHCullingPruneTrainer,
     "camera-densify-prune-shculling": CameraSHCullingPruningDensifyTrainer,
 }
+cameramodes = {k: lambda *args, **kwargs: ImportancePruningTrainerWrapper(v, *args, **kwargs) for k, v in cameramodes.items()}
 
 
 def prepare_quantizer(

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/reduced_3dgs.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reduced_3dgs
-Version: 1.8.14
+Version: 1.8.16
 Summary: Refactored code for the paper "Reducing the Memory Footprint of 3D Gaussian Splatting"
 Home-page: https://github.com/yindaheng98/reduced-3dgs
 Author: yindaheng98

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/reduced_3dgs.egg-info/SOURCES.txt

@@ -13,6 +13,8 @@ reduced_3dgs.egg-info/top_level.txt
 reduced_3dgs/pruning/__init__.py
 reduced_3dgs/pruning/combinations.py
 reduced_3dgs/pruning/trainer.py
+reduced_3dgs/pruning/importance/__init__.py
+reduced_3dgs/pruning/importance/trainer.py
 reduced_3dgs/quantization/__init__.py
 reduced_3dgs/quantization/abc.py
 reduced_3dgs/quantization/exclude_zeros.py
@@ -30,7 +32,13 @@ submodules/diff-gaussian-rasterization/cuda_rasterizer/rasterizer_impl.cu
 submodules/diff-gaussian-rasterization/reduced_3dgs/kmeans.cu
 submodules/diff-gaussian-rasterization/reduced_3dgs/redundancy_score.cu
 submodules/diff-gaussian-rasterization/reduced_3dgs/sh_culling.cu
+submodules/gaussian-importance/ext.cpp
+submodules/gaussian-importance/rasterize_points.cu
+submodules/gaussian-importance/cuda_rasterizer/backward.cu
+submodules/gaussian-importance/cuda_rasterizer/forward.cu
+submodules/gaussian-importance/cuda_rasterizer/rasterizer_impl.cu
 submodules/simple-knn/ext.cpp
 submodules/simple-knn/simple_knn.cu
 submodules/simple-knn/spatial.cu
-submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py
+submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/__init__.py
+submodules/gaussian-importance/diff_gaussian_rasterization/__init__.py

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/setup.py

@@ -40,6 +40,17 @@ rasterizor_packages = {
     'reduced_3dgs.simple_knn': 'submodules/simple-knn/simple_knn',
 }
 
+importance_root = "submodules/gaussian-importance"
+importance_sources = [
+    "cuda_rasterizer/rasterizer_impl.cu",
+    "cuda_rasterizer/forward.cu",
+    "cuda_rasterizer/backward.cu",
+    "rasterize_points.cu",
+    "ext.cpp"]
+importance_packages = {
+    'reduced_3dgs.pruning.importance.diff_gaussian_rasterization': 'submodules/gaussian-importance/diff_gaussian_rasterization',
+}
+
 cxx_compiler_flags = []
 nvcc_compiler_flags = []
 
@@ -49,7 +60,7 @@ if os.name == 'nt':
 
 setup(
     name="reduced_3dgs",
-    version='1.8.14',
+    version='1.8.16',
     author='yindaheng98',
     author_email='yindaheng98@gmail.com',
     url='https://github.com/yindaheng98/reduced-3dgs',
@@ -59,10 +70,11 @@ setup(
     classifiers=[
         "Programming Language :: Python :: 3",
     ],
-    packages=packages + list(rasterizor_packages.keys()),
+    packages=packages + list(rasterizor_packages.keys()) + list(importance_packages.keys()),
     package_dir={
         'reduced_3dgs': 'reduced_3dgs',
-        **rasterizor_packages
+        **rasterizor_packages,
+        **importance_packages,
     },
     ext_modules=[
         CUDAExtension(
@@ -70,6 +82,11 @@ setup(
             sources=[os.path.join(rasterizor_root, source) for source in rasterizor_sources],
             extra_compile_args={"nvcc": nvcc_compiler_flags + ["-I" + os.path.join(os.path.abspath(rasterizor_root), "third_party/glm/")]}
         ),
+        CUDAExtension(
+            name="reduced_3dgs.pruning.importance.diff_gaussian_rasterization._C",
+            sources=[os.path.join(importance_root, source) for source in importance_sources],
+            extra_compile_args={"nvcc": nvcc_compiler_flags + ["-I" + os.path.join(os.path.abspath(importance_root), "third_party/glm/")]}
+        ),
         CUDAExtension(
             name="reduced_3dgs.simple_knn._C",
             sources=[os.path.join(simpleknn_root, source) for source in simpleknn_sources],

{reduced_3dgs-1.8.14 → reduced_3dgs-1.8.16}/submodules/diff-gaussian-rasterization/reduced_3dgs/sh_culling.cu

@@ -1,8 +1,27 @@
 #include "sh_culling.h"
-#include "../cuda_rasterizer/auxiliary.h"
 #include <cooperative_groups.h>
 namespace cg = cooperative_groups;
 
+// Spherical harmonics coefficients
+__device__ const float SH_C0 = 0.28209479177387814f;
+__device__ const float SH_C1 = 0.4886025119029199f;
+__device__ const float SH_C2[] = {
+	1.0925484305920792f,
+	-1.0925484305920792f,
+	0.31539156525252005f,
+	-1.0925484305920792f,
+	0.5462742152960396f
+};
+__device__ const float SH_C3[] = {
+	-0.5900435899266435f,
+	2.890611442640554f,
+	-0.4570457994644658f,
+	0.3731763325901154f,
+	-0.4570457994644658f,
+	1.445305721320277f,
+	-0.5900435899266435f
+};
+
 __device__ void computeColorFromSH(const int idx, const int *degs, int max_coeffs, const glm::vec3 *means, glm::vec3 campos, const float *shs, glm::vec3 *out_colours)
 {
     // The implementation is loosely based on code for