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

reduced_3dgs/importance/diff_gaussian_rasterization/__init__.py

@@ -0,0 +1,347 @@
+#
+# 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,
+    colors_precomp,
+    opacities,
+    scales,
+    rotations,
+    cov3Ds_precomp,
+    raster_settings,
+):
+    if raster_settings.f_count:
+        return _RasterizeGaussians.forward_count(
+        means3D,
+        means2D,
+        sh,
+        colors_precomp,
+        opacities,
+        scales,
+        rotations,
+        cov3Ds_precomp,
+        raster_settings,
+
+    )
+    return _RasterizeGaussians.apply(
+    means3D,
+    means2D,
+    sh,
+    colors_precomp,
+    opacities,
+    scales,
+    rotations,
+    cov3Ds_precomp,
+    raster_settings,
+    )
+
+class _RasterizeGaussians(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        means3D,
+        means2D,
+        sh,
+        colors_precomp,
+        opacities,
+        scales,
+        rotations,
+        cov3Ds_precomp,
+        raster_settings,
+    ):
+
+        # 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,
+            raster_settings.sh_degree,
+            raster_settings.campos,
+            raster_settings.prefiltered,
+            raster_settings.debug
+        )
+        gaussians_count, opacity_important_score, T_alpha_important_score, num_rendered, color, radii, geomBuffer, \
+            binningBuffer, imgBuffer = None, None, None, None, None, None, None, None, None
+        
+        if raster_settings.f_count:
+            args = args + (raster_settings.f_count,)
+            if raster_settings.debug:
+                cpu_args = cpu_deep_copy_tuple(args)
+                try:
+                    gaussians_count, opacity_important_score, T_alpha_important_score, num_rendered, color, radii,\
+                        geomBuffer, binningBuffer, imgBuffer = _C.count_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:
+                gaussians_count, opacity_important_score, T_alpha_important_score, num_rendered, color, radii,\
+                    geomBuffer, binningBuffer, imgBuffer = _C.count_gaussians(*args)
+        else:
+            # 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)
+        ctx.count = gaussians_count
+        ctx.opacity_important_score = opacity_important_score
+        ctx.T_alpha_important_score = T_alpha_important_score
+        
+        if raster_settings.f_count: 
+            return gaussians_count, opacity_important_score, T_alpha_important_score, color, radii 
+        
+        return color, radii
+
+    @staticmethod
+    def forward_count(
+        means3D,
+        means2D,
+        sh,
+        colors_precomp,
+        opacities,
+        scales,
+        rotations,
+        cov3Ds_precomp,
+        raster_settings,
+    ):
+        assert(raster_settings.f_count)
+        # 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,
+            raster_settings.sh_degree,
+            raster_settings.campos,
+            raster_settings.prefiltered,
+            raster_settings.debug,
+            raster_settings.f_count
+        )
+        # gaussians_count, important_score, num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = None, None, None, None, None, None, None, None
+        # Invoke C++/CUDA rasterizer
+        # TODO(Kevin): pass the count in, but the output include a count list 
+        if raster_settings.debug:
+            cpu_args = cpu_deep_copy_tuple(args) # Copy them before they can be corrupted
+            try:
+                gaussians_count, opacity_important_score, T_alpha_important_score, num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.count_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:
+            gaussians_count, opacity_important_score, T_alpha_important_score, num_rendered, color, radii, geomBuffer, binningBuffer, imgBuffer = _C.count_gaussians(*args)
+                   
+        return gaussians_count, opacity_important_score, T_alpha_important_score, color, radii 
+
+    @staticmethod
+    def backward(ctx, grad_out_color, _):
+
+        # Restore necessary values from context
+        num_rendered = ctx.num_rendered
+        raster_settings = ctx.raster_settings
+        colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer = 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, 
+                raster_settings.sh_degree, 
+                raster_settings.campos,
+                geomBuffer,
+                num_rendered,
+                binningBuffer,
+                imgBuffer,
+                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,
+            grad_colors_precomp,
+            grad_opacities,
+            grad_scales,
+            grad_rotations,
+            grad_cov3Ds_precomp,
+            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
+    f_count : 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, colors_precomp = None, scales = None, rotations = None, cov3D_precomp = None):
+        
+        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,
+            colors_precomp,
+            opacities,
+            scales, 
+            rotations,
+            cov3D_precomp,
+            raster_settings, 
+        )
+    
+    def forward_count(self, means3D, means2D, opacities, shs = None, colors_precomp = None, scales = None, rotations = None, cov3D_precomp = None):
+        
+        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,
+            colors_precomp,
+            opacities,
+            scales, 
+            rotations,
+            cov3D_precomp,
+            raster_settings, 
+        ) 
+

reduced_3dgs/importance/trainer.py

@@ -0,0 +1,269 @@
+import math
+from typing import List
+import torch
+
+from gaussian_splatting import Camera, GaussianModel
+from gaussian_splatting.camera import build_camera
+from gaussian_splatting.trainer import AbstractDensifier, DensifierWrapper, DensificationTrainer, NoopDensifier
+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
+    }
+
+
+def prune_list(model: GaussianModel, dataset: CameraDataset, resize=None):
+    gaussian_count = torch.zeros(model.get_xyz.shape[0], device=model.get_xyz.device, dtype=torch.int)
+    opacity_important_score = torch.zeros(model.get_xyz.shape[0], device=model.get_xyz.device, dtype=torch.float)
+    T_alpha_important_score = torch.zeros(model.get_xyz.shape[0], device=model.get_xyz.device, dtype=torch.float)
+    for camera in dataset:
+        if resize is not None:
+            height, width = camera.image_height, camera.image_width
+            scale = resize / max(height, width)
+            height, width = int(height * scale), int(width * scale)
+            camera = build_camera(
+                image_height=height, image_width=width,
+                FoVx=camera.FoVx, FoVy=camera.FoVy,
+                R=camera.R, T=camera.T,
+                device=camera.R.device)
+        out = count_render(model, camera)
+        gaussian_count += out["gaussians_count"]
+        opacity_important_score += out["opacity_important_score"]
+        T_alpha_important_score += out["T_alpha_important_score"]
+    return gaussian_count, opacity_important_score, T_alpha_important_score
+
+
+# return importance score with adaptive volume measure described in paper
+def calculate_v_imp_score(gaussians: GaussianModel, imp_list, v_pow):
+    """
+    :param gaussians: A data structure containing Gaussian components with a get_scaling method.
+    :param imp_list: The importance scores for each Gaussian component.
+    :param v_pow: The power to which the volume ratios are raised.
+    :return: A list of adjusted values (v_list) used for pruning.
+    """
+    # Calculate the volume of each Gaussian component
+    volume = torch.prod(gaussians.get_scaling, dim=1)
+    # Determine the kth_percent_largest value
+    index = int(len(volume) * 0.9)
+    sorted_volume, _ = torch.sort(volume, descending=True)
+    kth_percent_largest = sorted_volume[index]
+    # Calculate v_list
+    v_list = torch.pow(volume / kth_percent_largest, v_pow)
+    v_list = v_list * imp_list
+    return v_list
+
+
+def score2mask(percent, import_score: list, threshold=None):
+    sorted_tensor, _ = torch.sort(import_score, dim=0)
+    index_nth_percentile = int(percent * (sorted_tensor.shape[0] - 1))
+    value_nth_percentile = sorted_tensor[index_nth_percentile]
+    thr = min(threshold, value_nth_percentile) if threshold is not None else value_nth_percentile
+    prune_mask = (import_score <= thr)
+    return prune_mask
+
+
+def prune_gaussians(
+        gaussians: GaussianModel, dataset: CameraDataset,
+        resize=None,
+        prune_type="comprehensive",
+        prune_percent=0.1,
+        prune_thr_important_score=None,
+        prune_thr_v_important_score=None,
+        prune_thr_max_v_important_score=None,
+        prune_thr_count=None,
+        prune_thr_T_alpha=None,
+        prune_thr_T_alpha_avg=None,
+        v_pow=0.1):
+    gaussian_list, opacity_imp_list, T_alpha_imp_list = prune_list(gaussians, dataset, resize)
+    match prune_type:
+        case "important_score":
+            mask = score2mask(prune_percent, opacity_imp_list, prune_thr_important_score)
+        case "v_important_score":
+            v_list = calculate_v_imp_score(gaussians, opacity_imp_list, v_pow)
+            mask = score2mask(prune_percent, v_list, prune_thr_v_important_score)
+        case "max_v_important_score":
+            v_list = opacity_imp_list * torch.max(gaussians.get_scaling, dim=1)[0]
+            mask = score2mask(prune_percent, v_list, prune_thr_max_v_important_score)
+        case "count":
+            mask = score2mask(prune_percent, gaussian_list, prune_thr_count)
+        case "T_alpha":
+            # new importance score defined by doji
+            mask = score2mask(prune_percent, T_alpha_imp_list, prune_thr_T_alpha)
+        case "T_alpha_avg":
+            v_list = T_alpha_imp_list / gaussian_list
+            v_list[gaussian_list <= 0] = 0
+            mask = score2mask(prune_percent, v_list, prune_thr_T_alpha_avg)
+        case "comprehensive":
+            mask = torch.zeros_like(gaussian_list, dtype=torch.bool)
+            if prune_thr_important_score is not None:
+                mask |= score2mask(prune_percent, opacity_imp_list, prune_thr_important_score)
+            if prune_thr_v_important_score is not None:
+                v_list = calculate_v_imp_score(gaussians, opacity_imp_list, v_pow)
+                mask |= score2mask(prune_percent, v_list, prune_thr_v_important_score)
+            if prune_thr_max_v_important_score is not None:
+                v_list = opacity_imp_list * torch.max(gaussians.get_scaling, dim=1)[0]
+                mask |= score2mask(prune_percent, v_list, prune_thr_max_v_important_score)
+            if prune_thr_count is not None:
+                mask |= score2mask(prune_percent, gaussian_list, prune_thr_count)
+            if prune_thr_T_alpha is not None:
+                mask |= score2mask(prune_percent, T_alpha_imp_list, prune_thr_T_alpha)
+            if prune_thr_T_alpha_avg is not None:
+                v_list = T_alpha_imp_list / gaussian_list
+                v_list[gaussian_list <= 0] = 0
+                mask |= score2mask(prune_percent, v_list, prune_thr_T_alpha_avg)
+        case _:
+            raise Exception("Unsupportive prunning method")
+    return mask
+
+
+class ImportancePruner(DensifierWrapper):
+    def __init__(
+            self, base_densifier: AbstractDensifier,
+            dataset: CameraDataset,
+            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,
+            importance_prune_thr_T_alpha_avg=0.001,
+            importance_v_pow=0.1):
+        super().__init__(base_densifier)
+        self.dataset = dataset
+        self.importance_prune_from_iter = importance_prune_from_iter
+        self.importance_prune_until_iter = importance_prune_until_iter
+        self.importance_prune_interval = importance_prune_interval
+        self.resize = importance_score_resize
+        self.prune_percent = importance_prune_percent
+        self.prune_thr_important_score = importance_prune_thr_important_score
+        self.prune_thr_v_important_score = importance_prune_thr_v_important_score
+        self.prune_thr_max_v_important_score = importance_prune_thr_max_v_important_score
+        self.prune_thr_count = importance_prune_thr_count
+        self.prune_thr_T_alpha = importance_prune_thr_T_alpha
+        self.prune_thr_T_alpha_avg = importance_prune_thr_T_alpha_avg
+        self.v_pow = importance_v_pow
+        self.prune_type = importance_prune_type
+
+    def densify_and_prune(self, loss, out, camera, step: int):
+        ret = super().densify_and_prune(loss, out, camera, step)
+        if self.importance_prune_from_iter <= step <= self.importance_prune_until_iter and step % self.importance_prune_interval == 0:
+            remove_mask = prune_gaussians(
+                self.model, self.dataset,
+                self.resize,
+                self.prune_type, self.prune_percent,
+                self.prune_thr_important_score, self.prune_thr_v_important_score,
+                self.prune_thr_max_v_important_score, self.prune_thr_count,
+                self.prune_thr_T_alpha, self.prune_thr_T_alpha_avg, self.v_pow,
+            )
+            ret = ret._replace(remove_mask=remove_mask if ret.remove_mask is None else torch.logical_or(ret.remove_mask, remove_mask))
+        return ret
+
+
+def BaseImportancePruningTrainer(
+        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 DensificationTrainer(
+        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,
+        ), *args, **kwargs
+    )

reduced_3dgs/pruning/__init__.py

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

reduced_3dgs/pruning/combinations.py

@@ -0,0 +1,65 @@
+
+from typing import Callable, List
+from gaussian_splatting import Camera, GaussianModel
+from gaussian_splatting.dataset import TrainableCameraDataset
+from gaussian_splatting.trainer import AbstractDensifier, DepthTrainerWrapper, NoopDensifier, SplitCloneDensifierTrainerWrapper
+from .trainer import BasePruner, BasePruningTrainer
+
+
+def PrunerInDensifyTrainerWrapper(
+        noargs_base_densifier_constructor: Callable[[GaussianModel, float, List[Camera]], AbstractDensifier],
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        *args,
+        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',
+        **kwargs):
+    return SplitCloneDensifierTrainerWrapper(
+        lambda model, scene_extent: BasePruner(
+            noargs_base_densifier_constructor(model, scene_extent, dataset),
+            scene_extent, 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
+    )
+
+
+def BasePrunerInDensifyTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        *args, **kwargs):
+    return PrunerInDensifyTrainerWrapper(
+        lambda model, scene_extent, dataset: NoopDensifier(model),
+        model, scene_extent, dataset,
+        *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