reduced-3dgs 1.10.0__cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl

reduced_3dgs/combinations.py

@@ -0,0 +1,245 @@
+from typing import List
+from gaussian_splatting import GaussianModel, CameraTrainableGaussianModel, Camera
+from gaussian_splatting.dataset import CameraDataset, TrainableCameraDataset
+from gaussian_splatting.trainer import OpacityResetDensificationTrainer
+# from gaussian_splatting.trainer import BaseOpacityResetDensificationTrainer as OpacityResetDensificationTrainer
+from gaussian_splatting.trainer import OpacityResetTrainerWrapper, CameraTrainerWrapper, NoopDensifier, DepthTrainerWrapper
+from .shculling import VariableSHGaussianModel, SHCullingTrainerWrapper
+from .shculling import SHCullingTrainer
+# from .shculling import BaseSHCullingTrainer as SHCullingTrainer
+from .pruning import PruningTrainerWrapper, PrunerInDensifyTrainerWrapper
+# from .pruning import BasePruningTrainer as PruningTrainer, BasePrunerInDensifyTrainer as PrunerInDensifyTrainer
+from .importance import ImportancePruner
+
+
+def BaseFullPruningTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        *args,
+        importance_prune_from_iter=15000,
+        importance_prune_until_iter=20000,
+        importance_prune_interval: int = 1000,
+        importance_score_resize=None,
+        importance_prune_type="comprehensive",
+        importance_prune_percent=0.1,
+        importance_prune_thr_important_score=None,
+        importance_prune_thr_v_important_score=3.0,
+        importance_prune_thr_max_v_important_score=None,
+        importance_prune_thr_count=1,
+        importance_prune_thr_T_alpha=1.0,
+        importance_prune_thr_T_alpha_avg=0.001,
+        importance_v_pow=0.1,
+        **kwargs):
+    return PruningTrainerWrapper(
+        lambda model, scene_extent, dataset: ImportancePruner(
+            NoopDensifier(model),
+            dataset,
+            importance_prune_from_iter=importance_prune_from_iter,
+            importance_prune_until_iter=importance_prune_until_iter,
+            importance_prune_interval=importance_prune_interval,
+            importance_score_resize=importance_score_resize,
+            importance_prune_type=importance_prune_type,
+            importance_prune_percent=importance_prune_percent,
+            importance_prune_thr_important_score=importance_prune_thr_important_score,
+            importance_prune_thr_v_important_score=importance_prune_thr_v_important_score,
+            importance_prune_thr_max_v_important_score=importance_prune_thr_max_v_important_score,
+            importance_prune_thr_count=importance_prune_thr_count,
+            importance_prune_thr_T_alpha=importance_prune_thr_T_alpha,
+            importance_prune_thr_T_alpha_avg=importance_prune_thr_T_alpha_avg,
+            importance_v_pow=importance_v_pow,
+        ),
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def BaseFullPrunerInDensifyTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        *args,
+        importance_prune_from_iter=15000,
+        importance_prune_until_iter=20000,
+        importance_prune_interval: int = 1000,
+        importance_score_resize=None,
+        importance_prune_type="comprehensive",
+        importance_prune_percent=0.1,
+        importance_prune_thr_important_score=None,
+        importance_prune_thr_v_important_score=3.0,
+        importance_prune_thr_max_v_important_score=None,
+        importance_prune_thr_count=1,
+        importance_prune_thr_T_alpha=1.0,
+        importance_prune_thr_T_alpha_avg=0.001,
+        importance_v_pow=0.1,
+        **kwargs):
+    return PrunerInDensifyTrainerWrapper(
+        lambda model, scene_extent, dataset: ImportancePruner(
+            NoopDensifier(model),
+            dataset,
+            importance_prune_from_iter=importance_prune_from_iter,
+            importance_prune_until_iter=importance_prune_until_iter,
+            importance_prune_interval=importance_prune_interval,
+            importance_score_resize=importance_score_resize,
+            importance_prune_type=importance_prune_type,
+            importance_prune_percent=importance_prune_percent,
+            importance_prune_thr_important_score=importance_prune_thr_important_score,
+            importance_prune_thr_v_important_score=importance_prune_thr_v_important_score,
+            importance_prune_thr_max_v_important_score=importance_prune_thr_max_v_important_score,
+            importance_prune_thr_count=importance_prune_thr_count,
+            importance_prune_thr_T_alpha=importance_prune_thr_T_alpha,
+            importance_prune_thr_T_alpha_avg=importance_prune_thr_T_alpha_avg,
+            importance_v_pow=importance_v_pow,
+        ),
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def DepthFullPruningTrainer(model: GaussianModel, scene_extent: float, dataset: TrainableCameraDataset, *args, **kwargs):
+    return DepthTrainerWrapper(BaseFullPruningTrainer, model, scene_extent, *args, dataset=dataset, **kwargs)
+
+
+def DepthFullPrunerInDensifyTrainer(model: GaussianModel, scene_extent: float, dataset: TrainableCameraDataset, *args, **kwargs):
+    return DepthTrainerWrapper(BaseFullPrunerInDensifyTrainer, model, scene_extent, *args, dataset=dataset, **kwargs)
+
+
+def OpacityResetPruningTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        *args, **kwargs):
+    return OpacityResetTrainerWrapper(
+        lambda model, scene_extent, *args, **kwargs: DepthFullPruningTrainer(model, scene_extent, dataset, *args, **kwargs),
+        model, scene_extent,
+        *args, **kwargs
+    )
+
+
+def OpacityResetPrunerInDensifyTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        *args, **kwargs):
+    return OpacityResetTrainerWrapper(
+        lambda model, scene_extent, *args, **kwargs: DepthFullPrunerInDensifyTrainer(model, scene_extent, dataset, *args, **kwargs),
+        model, scene_extent,
+        *args, **kwargs
+    )
+
+
+PruningTrainer = OpacityResetPruningTrainer
+PrunerInDensifyTrainer = OpacityResetPrunerInDensifyTrainer
+
+
+def SHCullingDensificationTrainer(
+    model: VariableSHGaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        *args, **kwargs):
+    return SHCullingTrainerWrapper(
+        lambda model, scene_extent, dataset, *args, **kwargs: OpacityResetDensificationTrainer(model, scene_extent, *args, **kwargs),
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def SHCullingPruningTrainer(
+    model: VariableSHGaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        *args, **kwargs):
+    return SHCullingTrainerWrapper(
+        OpacityResetPruningTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def SHCullingPrunerInDensifyTrainer(
+    model: VariableSHGaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        *args, **kwargs):
+    return SHCullingTrainerWrapper(
+        OpacityResetPrunerInDensifyTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+class CameraTrainableVariableSHGaussianModel(VariableSHGaussianModel):
+    def forward(self, camera: Camera):
+        return CameraTrainableGaussianModel.forward(self, camera)
+
+
+def CameraSHCullingTrainer(
+        model: CameraTrainableVariableSHGaussianModel,
+        scene_extent: float,
+        dataset: TrainableCameraDataset,
+        *args, **kwargs):
+    return CameraTrainerWrapper(
+        SHCullingTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def CameraPruningTrainer(
+        model: CameraTrainableVariableSHGaussianModel,
+        scene_extent: float,
+        dataset: TrainableCameraDataset,
+        *args, **kwargs):
+    return CameraTrainerWrapper(
+        OpacityResetPruningTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def CameraPrunerInDensifyTrainer(
+        model: CameraTrainableVariableSHGaussianModel,
+        scene_extent: float,
+        dataset: TrainableCameraDataset,
+        *args, **kwargs):
+    return CameraTrainerWrapper(
+        OpacityResetPrunerInDensifyTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def CameraSHCullingDensifyTrainer(
+        model: CameraTrainableVariableSHGaussianModel,
+        scene_extent: float,
+        dataset: TrainableCameraDataset,
+        *args, **kwargs):
+    return CameraTrainerWrapper(
+        SHCullingDensificationTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def CameraSHCullingPruningTrainer(
+        model: CameraTrainableVariableSHGaussianModel,
+        scene_extent: float,
+        dataset: TrainableCameraDataset,
+        *args, **kwargs):
+    return CameraTrainerWrapper(
+        SHCullingPruningTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )
+
+
+def CameraSHCullingPruningDensifyTrainer(
+        model: CameraTrainableVariableSHGaussianModel,
+        scene_extent: float,
+        dataset: TrainableCameraDataset,
+        *args, **kwargs):
+    return CameraTrainerWrapper(
+        SHCullingPrunerInDensifyTrainer,
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )

reduced_3dgs/diff_gaussian_rasterization/__init__.py

@@ -0,0 +1,235 @@
+#
+# Copyright (C) 2023, Inria
+# GRAPHDECO research group, https://team.inria.fr/graphdeco
+# All rights reserved.
+#
+# This software is free for non-commercial, research and evaluation use 
+# under the terms of the LICENSE.md file.
+#
+# For inquiries contact  george.drettakis@inria.fr
+#
+
+from typing import NamedTuple
+import torch.nn as nn
+import torch
+from . import _C
+
+def cpu_deep_copy_tuple(input_tuple):
+    copied_tensors = [item.cpu().clone() if isinstance(item, torch.Tensor) else item for item in input_tuple]
+    return tuple(copied_tensors)
+
+def rasterize_gaussians(
+    means3D,
+    means2D,
+    sh,
+    degrees,
+    colors_precomp,
+    opacities,
+    scales,
+    rotations,
+    cov3Ds_precomp,
+    raster_settings,
+    lambda_sh_sparsity
+):
+    return _RasterizeGaussians.apply(
+        means3D,
+        means2D,
+        sh,
+        degrees,
+        colors_precomp,
+        opacities,
+        scales,
+        rotations,
+        cov3Ds_precomp,
+        raster_settings,
+        lambda_sh_sparsity
+    )
+
+class _RasterizeGaussians(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        means3D,
+        means2D,
+        sh,
+        degrees,
+        colors_precomp,
+        opacities,
+        scales,
+        rotations,
+        cov3Ds_precomp,
+        raster_settings,
+        lambda_sh_sparsity
+    ):
+
+        # Restructure arguments the way that the C++ lib expects them
+        args = (
+            raster_settings.bg, 
+            means3D,
+            colors_precomp,
+            opacities,
+            scales,
+            rotations,
+            raster_settings.scale_modifier,
+            cov3Ds_precomp,
+            raster_settings.viewmatrix,
+            raster_settings.projmatrix,
+            raster_settings.tanfovx,
+            raster_settings.tanfovy,
+            raster_settings.image_height,
+            raster_settings.image_width,
+            sh,
+            degrees,
+            raster_settings.campos,
+            raster_settings.prefiltered,
+            True
+            # raster_settings.debug
+        )
+
+        # Invoke C++/CUDA rasterizer
+        if raster_settings.debug:
+            cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
+            try:
+                num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args)
+            except Exception as ex:
+                torch.save(cpu_args, "snapshot_fw.dump")
+                print("\nAn error occured in forward. Please forward snapshot_fw.dump for debugging.")
+                raise ex
+        else:
+            num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.rasterize_gaussians(*args)
+
+        # Keep relevant tensors for backward
+        ctx.raster_settings = raster_settings
+        ctx.num_rendered = num_rendered
+        ctx.save_for_backward(colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer, degrees)
+        ctx.lambda_sh_sparsity = lambda_sh_sparsity
+        return color, radii
+
+    @staticmethod
+    def backward(ctx, grad_out_color, _):
+
+        # Restore necessary values from context
+        num_rendered = ctx.num_rendered
+        raster_settings = ctx.raster_settings
+        lambda_sh_sparsity = ctx.lambda_sh_sparsity
+        colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer, degrees = ctx.saved_tensors
+
+        # Restructure args as C++ method expects them
+        args = (raster_settings.bg,
+                means3D, 
+                radii, 
+                colors_precomp, 
+                scales, 
+                rotations, 
+                raster_settings.scale_modifier, 
+                cov3Ds_precomp, 
+                raster_settings.viewmatrix, 
+                raster_settings.projmatrix, 
+                raster_settings.tanfovx, 
+                raster_settings.tanfovy, 
+                grad_out_color,
+                sh,
+                degrees, 
+                raster_settings.campos,
+                geomBuffer,
+                num_rendered,
+                binningBuffer,
+                imgBuffer,
+                lambda_sh_sparsity,
+                raster_settings.debug)
+
+        # Compute gradients for relevant tensors by invoking backward method
+        if raster_settings.debug:
+            cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
+            try:
+                grad_means2D, grad_colors_precomp, grad_opacities, grad_means3D, grad_cov3Ds_precomp, grad_sh, grad_scales, grad_rotations = _C.rasterize_gaussians_backward(*args)
+            except Exception as ex:
+                torch.save(cpu_args, "snapshot_bw.dump")
+                print("\nAn error occured in backward. Writing snapshot_bw.dump for debugging.\n")
+                raise ex
+        else:
+            grad_means2D, grad_colors_precomp, grad_opacities, grad_means3D, grad_cov3Ds_precomp, grad_sh, grad_scales, grad_rotations = _C.rasterize_gaussians_backward(*args)
+
+        grads = (
+            grad_means3D,
+            grad_means2D,
+            grad_sh,
+            None,
+            grad_colors_precomp,
+            grad_opacities,
+            grad_scales,
+            grad_rotations,
+            grad_cov3Ds_precomp,
+            None,
+            None,
+        )
+
+        return grads
+
+class GaussianRasterizationSettings(NamedTuple):
+    image_height: int
+    image_width: int 
+    tanfovx : float
+    tanfovy : float
+    bg : torch.Tensor
+    scale_modifier : float
+    viewmatrix : torch.Tensor
+    projmatrix : torch.Tensor
+    sh_degree : int
+    campos : torch.Tensor
+    prefiltered : bool
+    debug : bool
+
+class GaussianRasterizer(nn.Module):
+    def __init__(self, raster_settings):
+        super().__init__()
+        self.raster_settings = raster_settings
+
+    def markVisible(self, positions):
+        # Mark visible points (based on frustum culling for camera) with a boolean 
+        with torch.no_grad():
+            raster_settings = self.raster_settings
+            visible = _C.mark_visible(
+                positions,
+                raster_settings.viewmatrix,
+                raster_settings.projmatrix)
+            
+        return visible
+
+    def forward(self, means3D, means2D, opacities, shs = None, degrees = None, colors_precomp = None, scales = None, rotations = None, cov3D_precomp = None, lambda_sh_sparsity=0.):
+        
+        raster_settings = self.raster_settings
+
+        if (shs is None and colors_precomp is None) or (shs is not None and colors_precomp is not None):
+            raise Exception('Please provide excatly one of either SHs or precomputed colors!')
+        
+        if ((scales is None or rotations is None) and cov3D_precomp is None) or ((scales is not None or rotations is not None) and cov3D_precomp is not None):
+            raise Exception('Please provide exactly one of either scale/rotation pair or precomputed 3D covariance!')
+        
+        if shs is None:
+            shs = torch.Tensor([])
+        if colors_precomp is None:
+            colors_precomp = torch.Tensor([])
+
+        if scales is None:
+            scales = torch.Tensor([])
+        if rotations is None:
+            rotations = torch.Tensor([])
+        if cov3D_precomp is None:
+            cov3D_precomp = torch.Tensor([])
+
+        # Invoke C++/CUDA rasterization routine
+        return rasterize_gaussians(
+            means3D,
+            means2D,
+            shs,
+            degrees,
+            colors_precomp,
+            opacities,
+            scales, 
+            rotations,
+            cov3D_precomp,
+            raster_settings, 
+            lambda_sh_sparsity
+        )
+

reduced_3dgs/importance/__init__.py

@@ -0,0 +1,3 @@
+from .trainer import ImportancePruner, BaseImportancePruningTrainer
+from .combinations import BaseImportancePrunerInDensifyTrainer, DepthImportancePruningTrainer, DepthImportancePrunerInDensifyTrainer
+from .combinations import ImportancePruningTrainer, ImportancePrunerInDensifyTrainer

reduced_3dgs/importance/combinations.py

@@ -0,0 +1,63 @@
+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 ImportancePruner, BaseImportancePruningTrainer
+
+
+def BaseImportancePrunerInDensifyTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        *args,
+        importance_prune_from_iter=15000,
+        importance_prune_until_iter=20000,
+        importance_prune_interval: int = 1000,
+        importance_score_resize=None,
+        importance_prune_type="comprehensive",
+        importance_prune_percent=0.1,
+        importance_prune_thr_important_score=None,
+        importance_prune_thr_v_important_score=3.0,
+        importance_prune_thr_max_v_important_score=None,
+        importance_prune_thr_count=1,
+        importance_prune_thr_T_alpha=1.0,
+        importance_prune_thr_T_alpha_avg=0.001,
+        importance_v_pow=0.1,
+        **kwargs):
+    return DensificationTrainerWrapper(
+        lambda model, scene_extent: ImportancePruner(
+            NoopDensifier(model),
+            dataset,
+            importance_prune_from_iter=importance_prune_from_iter,
+            importance_prune_until_iter=importance_prune_until_iter,
+            importance_prune_interval=importance_prune_interval,
+            importance_score_resize=importance_score_resize,
+            importance_prune_type=importance_prune_type,
+            importance_prune_percent=importance_prune_percent,
+            importance_prune_thr_important_score=importance_prune_thr_important_score,
+            importance_prune_thr_v_important_score=importance_prune_thr_v_important_score,
+            importance_prune_thr_max_v_important_score=importance_prune_thr_max_v_important_score,
+            importance_prune_thr_count=importance_prune_thr_count,
+            importance_prune_thr_T_alpha=importance_prune_thr_T_alpha,
+            importance_prune_thr_T_alpha_avg=importance_prune_thr_T_alpha_avg,
+            importance_v_pow=importance_v_pow,
+        ),
+        model,
+        scene_extent,
+        *args, **kwargs
+    )
+
+
+# Depth trainer
+
+
+def DepthImportancePruningTrainer(model: GaussianModel, scene_extent: float, dataset: TrainableCameraDataset, *args, **kwargs):
+    return DepthTrainerWrapper(BaseImportancePruningTrainer, model, scene_extent, *args, dataset=dataset, **kwargs)
+
+
+def DepthImportancePrunerInDensifyTrainer(model: GaussianModel, scene_extent: float, dataset: TrainableCameraDataset, *args, **kwargs):
+    return DepthTrainerWrapper(BaseImportancePrunerInDensifyTrainer, model, scene_extent, *args, dataset=dataset, **kwargs)
+
+
+ImportancePruningTrainer = DepthImportancePruningTrainer
+ImportancePrunerInDensifyTrainer = DepthImportancePrunerInDensifyTrainer