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

reduced_3dgs/pruning/trainer.py

@@ -0,0 +1,145 @@
+from typing import Callable, List
+import torch
+from gaussian_splatting import GaussianModel, Camera
+from gaussian_splatting.trainer import AbstractDensifier, OpacityPruner, 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
+
+
+def calculate_redundancy_metric(gaussians: GaussianModel, cameras: List[Camera], pixel_scale=1.0, num_neighbours=30):
+    # Get minimum projected pixel size
+    cube_size = find_minimum_projected_pixel_size(
+        torch.stack([camera.full_proj_transform for camera in cameras], dim=0),
+        torch.stack([camera.full_proj_transform.inverse() for camera in cameras], dim=0),
+        gaussians._xyz,
+        torch.tensor([camera.image_height for camera in cameras], device="cuda", dtype=torch.int32),
+        torch.tensor([camera.image_width for camera in cameras], device="cuda", dtype=torch.int32)
+    )
+
+    scaled_pixel_size = cube_size * pixel_scale
+    half_diagonal = scaled_pixel_size * torch.sqrt(torch.tensor([3], device="cuda")) / 2
+
+    # Find neighbours as candidates for the intersection test
+    _, indices = distIndex2(gaussians.get_xyz, num_neighbours)
+    indices = indices.view(-1, num_neighbours)
+
+    # Do the intersection check
+    redundancy_metrics, intersection_mask = sphere_ellipsoid_intersection(gaussians._xyz,
+                                                                          gaussians.get_scaling,
+                                                                          gaussians.get_rotation,
+                                                                          indices,
+                                                                          half_diagonal,
+                                                                          num_neighbours)
+    # We haven't counted count for the primitive at the center of each sphere, so add 1 to everything
+    redundancy_metrics += 1
+
+    indices = torch.cat((torch.arange(gaussians.get_xyz.shape[0], device="cuda", dtype=torch.int).view(-1, 1), indices), dim=1)
+    intersection_mask = torch.cat((torch.ones_like(gaussians._opacity, device="cuda", dtype=bool), intersection_mask), dim=1)
+
+    min_redundancy_metrics = allocate_minimum_redundancy_value(redundancy_metrics, indices, intersection_mask, num_neighbours+1)[0]
+    return min_redundancy_metrics, cube_size
+
+
+def mercy_points(self: GaussianModel, _splatted_num_accum: torch.Tensor, lambda_mercy=2, mercy_minimum=2, mercy_type='redundancy_opacity'):
+    mean = _splatted_num_accum.float().mean(dim=0, keepdim=True)
+    std = _splatted_num_accum.float().var(dim=0, keepdim=True).sqrt()
+
+    threshold = max((mean + lambda_mercy*std).item(), mercy_minimum)
+
+    mask = (_splatted_num_accum > threshold)
+
+    if mercy_type == 'redundancy_opacity':
+        # Prune redundant points based on 50% lowest opacity
+        mask[mask.clone()] = self.get_opacity[mask].squeeze() < self.get_opacity[mask].median()
+    elif mercy_type == 'redundancy_random':
+        # Prune 50% redundant points at random
+        mask[mask.clone()] = torch.rand(mask[mask].shape, device="cuda").squeeze() < 0.5
+    elif mercy_type == 'opacity':
+        # Prune based just on opacity
+        threshold = self.get_opacity.quantile(0.045)
+        mask = (self.get_opacity < threshold).squeeze()
+    elif mercy_type == 'redundancy_opacity_opacity':
+        # Prune based on opacity and on redundancy + opacity (options 1 and 3)
+        mask[mask.clone()] = self.get_opacity[mask].squeeze() < self.get_opacity[mask].median()
+        threshold = torch.min(self.get_opacity.quantile(0.03), torch.tensor([0.05], device="cuda"))
+        mask = torch.logical_or(mask, (self.get_opacity < threshold).squeeze())
+    return mask
+
+
+def mercy_gaussians(
+    model: GaussianModel,
+    dataset: List[Camera],
+    box_size=1.,
+    lambda_mercy=1.,
+    mercy_minimum=3,
+    mercy_type='redundancy_opacity'
+):
+    _splatted_num_accum, _ = calculate_redundancy_metric(model, dataset, pixel_scale=box_size)
+    mask = mercy_points(model, _splatted_num_accum.squeeze(), lambda_mercy, mercy_minimum, mercy_type)
+    return mask
+
+
+class BasePruner(OpacityPruner):
+    def __init__(
+            self, base_densifier: AbstractDensifier,
+            scene_extent,
+            dataset: List[Camera],
+            *args,
+            box_size=1.,
+            lambda_mercy=1.,
+            mercy_minimum=3,
+            mercy_type='redundancy_opacity',
+            **kwargs):
+        super().__init__(base_densifier, 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:
+        remove_mask = mercy_gaussians(self.model, self.dataset, self.box_size, self.lambda_mercy, self.mercy_minimum, self.mercy_type)
+        prune_mask = torch.logical_or(super().prune(), remove_mask)
+        return prune_mask
+
+
+def PruningTrainerWrapper(
+        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 DensificationTrainer(
+        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,
+        ), *args, **kwargs
+    )
+
+
+def BasePruningTrainer(
+        model: GaussianModel,
+        scene_extent: float,
+        dataset: List[Camera],
+        *args, **kwargs):
+    return DensificationTrainer(
+        lambda model, scene_extent, dataset: NoopDensifier(model),
+        model, scene_extent, dataset,
+        *args, **kwargs
+    )

reduced_3dgs/quantization/__init__.py

@@ -0,0 +1,4 @@
+from .abc import AbstractQuantizer, QuantizeTrainerWrapper
+from .quantizer import VectorQuantizer
+from .exclude_zeros import ExcludeZeroSHQuantizer
+from .wrapper import VectorQuantizeTrainerWrapper, VectorQuantizeTrainer

reduced_3dgs/quantization/abc.py

@@ -0,0 +1,49 @@
+import abc
+from typing import Dict, Tuple
+
+import torch
+
+from gaussian_splatting import GaussianModel
+from gaussian_splatting.trainer import AbstractTrainer, TrainerWrapper
+
+
+class AbstractQuantizer(abc.ABC):
+
+    @abc.abstractmethod
+    def quantize(self, model: GaussianModel, update_codebook=True) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        pass
+
+    @abc.abstractmethod
+    def dequantize(self, model: GaussianModel, ids_dict: Dict[str, torch.Tensor], codebook_dict: Dict[str, torch.Tensor], xyz: torch.Tensor = None, replace=False) -> GaussianModel:
+        pass
+
+    @abc.abstractmethod
+    def save_quantized(self, model: GaussianModel, ply_path: str, codebook_path: str):
+        pass
+
+    @abc.abstractmethod
+    def load_quantized(self, model: GaussianModel, ply_path: str, codebook_path: str) -> GaussianModel:
+        pass
+
+
+class QuantizeTrainerWrapper(TrainerWrapper, metaclass=abc.ABCMeta):
+    def __init__(
+            self, base_trainer: AbstractTrainer,
+            quantizer: AbstractQuantizer,
+            quantize_from_iter=5000,
+            quantize_until_iter=30000,
+            quantize_interval=1000,
+    ):
+        super().__init__(base_trainer)
+        self.quantizer = quantizer
+        self.quantize_from_iter = quantize_from_iter
+        self.quantize_until_iter = quantize_until_iter
+        self.quantize_interval = quantize_interval
+
+    @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:
+            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)
+        return self.base_trainer.model

reduced_3dgs/quantization/exclude_zeros.py

@@ -0,0 +1,41 @@
+import torch
+from gaussian_splatting import GaussianModel
+from .quantizer import VectorQuantizer
+
+
+class ExcludeZeroSHQuantizer(VectorQuantizer):
+    def __init__(self, *args, treat_as_zero=1e-8, **kwargs):
+        super(ExcludeZeroSHQuantizer, self).__init__(*args, **kwargs)
+        self.treat_as_zero = treat_as_zero
+
+    def zeros_mask(self, values: torch.Tensor):
+        return (values.abs() < self.treat_as_zero).all(-1)
+
+    def generate_codebook_exclude_zero(self, values: torch.Tensor, num_clusters=256, init_codebook=None):
+        zeros_mask = self.zeros_mask(values)
+        if zeros_mask.all():
+            return torch.zeros(1, values.shape[1], dtype=values.dtype, device=values.device), torch.zeros(values.shape[0], dtype=torch.long, device=values.device)
+        if init_codebook is not None:
+            if init_codebook.abs().max() < self.treat_as_zero:
+                init_codebook = None
+            elif init_codebook.shape[0] > num_clusters - 1:
+                init_codebook = init_codebook[-(num_clusters - 1):, ...]
+        nonzero_values = values[~zeros_mask]
+        nonzero_centers, nonzero_ids = super().generate_codebook(nonzero_values, num_clusters - 1, init_codebook)
+        ids = torch.zeros(values.shape[0], dtype=nonzero_ids.dtype, device=nonzero_ids.device)
+        ids[~zeros_mask] = nonzero_ids + 1
+        centers = torch.cat((torch.zeros(1, values.shape[1], dtype=values.dtype, device=values.device), nonzero_centers), dim=0)
+        return centers, ids
+
+    def has_zero(self, values: torch.Tensor):
+        return self.zeros_mask(values).any()
+
+    def produce_clusters_degree_features_rest(self, model: GaussianModel, sh_degree, *args, **kwargs):
+        features_rest_flatten = model._features_rest.detach().transpose(1, 2).flatten(0, 1)
+        sh_idx_start, sh_idx_end = (sh_degree + 1) ** 2 - 1, (sh_degree + 2) ** 2 - 1
+        features_rest = features_rest_flatten[:, sh_idx_start:sh_idx_end]
+        if self.has_zero(features_rest):
+            codebook, ids = self.generate_codebook_exclude_zero(features_rest, self.num_clusters_features_rest[sh_degree], *args, **kwargs)
+        else:
+            codebook, ids = self.generate_codebook(features_rest, self.num_clusters_features_rest[sh_degree], *args, **kwargs)
+        return codebook, ids.reshape(-1, model._features_rest.shape[-1])

reduced_3dgs/quantization/quantizer.py

@@ -0,0 +1,289 @@
+import math
+import os
+from typing import Dict, Tuple
+import torch
+import torch.nn as nn
+import numpy as np
+from sklearn.cluster import MiniBatchKMeans as KMeans
+from gaussian_splatting import GaussianModel
+from plyfile import PlyData, PlyElement
+import numpy as np
+from .abc import AbstractQuantizer
+
+
+def array2record(array: torch.Tensor, perfix, n_cols, dtype):
+    dtype_full = [(f'{perfix}_{i}', dtype) for i in range(n_cols)] if n_cols > 1 else [(perfix, dtype)]
+    data_full = map(lambda x: x.squeeze(-1), np.array_split(array.cpu().numpy(), n_cols, axis=1))
+    record = np.rec.fromarrays(data_full, dtype=dtype_full)
+    return record
+
+
+def compute_uint_length(n):
+    count = 0
+    while n >> 1:
+        count += 1
+        n >>= 1
+    return count
+
+
+def compute_uint_dtype(n):
+    bits = compute_uint_length(n)
+    bytes = bits // 8
+    if bits % 8:
+        bytes += 1
+    return f'u{bytes}'
+
+
+class VectorQuantizer(AbstractQuantizer):
+    def __init__(
+            self,
+            num_clusters=256,
+            num_clusters_rotation_re=None,
+            num_clusters_rotation_im=None,
+            num_clusters_opacity=None,
+            num_clusters_scaling=None,
+            num_clusters_features_dc=None,
+            num_clusters_features_rest=[],
+            max_sh_degree=3,
+            force_code_dtype=None,
+            force_codebook_dtype='f4',
+            tol=1e-6, max_iter=500,
+    ):
+        self.num_clusters_rotation_re = num_clusters_rotation_re or num_clusters
+        self.num_clusters_rotation_im = num_clusters_rotation_im or num_clusters
+        self.num_clusters_opacity = num_clusters_opacity or num_clusters
+        self.num_clusters_scaling = num_clusters_scaling or num_clusters
+        self.num_clusters_features_dc = num_clusters_features_dc or num_clusters
+        self.num_clusters_features_rest = [(num_clusters_features_rest[i] if len(num_clusters_features_rest) > i else num_clusters) for i in range(max_sh_degree)]
+        self.force_code_dtype = force_code_dtype
+        self.force_codebook_dtype = force_codebook_dtype
+        self.tol = tol
+        self.max_iter = max_iter
+
+        self._codebook_dict = {}
+
+    def generate_codebook(self, values: torch.Tensor, num_clusters, init_codebook=None):
+        kmeans = KMeans(
+            n_clusters=num_clusters, tol=self.tol, max_iter=self.max_iter,
+            init='k-means++' if init_codebook is None else init_codebook.cpu().numpy(),
+            random_state=0, n_init="auto", verbose=0,
+            batch_size=256 * os.cpu_count()
+        )
+        ids = torch.tensor(kmeans.fit_predict(values.cpu().numpy()), device=values.device)
+        centers = torch.tensor(kmeans.cluster_centers_, dtype=values.dtype, device=values.device)
+        return centers, ids
+
+    def one_nearst(self, points: torch.Tensor, codebook: torch.Tensor, batch=2**16):
+        ids = torch.zeros(points.shape[0], dtype=torch.int64, device=points.device)
+        for i in range(math.ceil(points.shape[0]/batch)):
+            ids[i*batch:i*batch+batch] = torch.argmin(torch.cdist(points[i*batch:i*batch+batch, ...], codebook), dim=1)
+        return ids
+
+    def produce_clusters_features_dc(self, model: GaussianModel, *args, **kwargs):
+        codebook, ids = self.generate_codebook(model._features_dc.detach().squeeze(1), self.num_clusters_features_dc, *args, **kwargs)
+        return codebook, ids.unsqueeze(1)
+
+    def find_nearest_cluster_id_features_dc(self, model: GaussianModel, codebook: torch.Tensor):
+        return self.one_nearst(model._features_dc.detach().squeeze(1), codebook).unsqueeze(1)
+
+    def produce_clusters_degree_features_rest(self, model: GaussianModel, sh_degree, *args, **kwargs):
+        features_rest_flatten = model._features_rest.detach().transpose(1, 2).flatten(0, 1)
+        sh_idx_start, sh_idx_end = (sh_degree + 1) ** 2 - 1, (sh_degree + 2) ** 2 - 1
+        features_rest = features_rest_flatten[:, sh_idx_start:sh_idx_end]
+        codebook, ids = self.generate_codebook(features_rest, self.num_clusters_features_rest[sh_degree], *args, **kwargs)
+        return codebook, ids.reshape(-1, model._features_rest.shape[-1])
+
+    def find_nearest_cluster_id_degree_features_rest(self, model: GaussianModel, sh_degree, codebook: torch.Tensor):
+        features_rest_flatten = model._features_rest.detach().transpose(1, 2).flatten(0, 1)
+        sh_idx_start, sh_idx_end = (sh_degree + 1) ** 2 - 1, (sh_degree + 2) ** 2 - 1
+        features_rest = features_rest_flatten[:, sh_idx_start:sh_idx_end]
+        ids = self.one_nearst(features_rest, codebook)
+        return ids.reshape(-1, model._features_rest.shape[-1])
+
+    def produce_clusters_rotation_re(self, model: GaussianModel, *args, **kwargs):
+        return self.generate_codebook(model.get_rotation.detach()[:, 0:1], self.num_clusters_rotation_re, *args, **kwargs)
+
+    def find_nearest_cluster_id_rotation_re(self, model: GaussianModel, codebook: torch.Tensor):
+        return self.one_nearst(model.get_rotation.detach()[:, 0:1], codebook)
+
+    def produce_clusters_rotation_im(self, model: GaussianModel, *args, **kwargs):
+        return self.generate_codebook(model.get_rotation.detach()[:, 1:], self.num_clusters_rotation_im, *args, **kwargs)
+
+    def find_nearest_cluster_id_rotation_im(self, model: GaussianModel, codebook: torch.Tensor):
+        return self.one_nearst(model.get_rotation.detach()[:, 1:], codebook)
+
+    def produce_clusters_opacity(self, model: GaussianModel, *args, **kwargs):
+        return self.generate_codebook(model._opacity.detach(), self.num_clusters_opacity, *args, **kwargs)
+
+    def find_nearest_cluster_id_opacity(self, model: GaussianModel, codebook: torch.Tensor):
+        return self.one_nearst(model._opacity.detach(), codebook)
+
+    def produce_clusters_scaling(self, model: GaussianModel, *args, **kwargs):
+        centers, ids = self.generate_codebook(model.get_scaling.detach(), self.num_clusters_scaling, *args, **kwargs)
+        centers_log = model.scaling_inverse_activation(centers)
+        return centers_log, ids
+
+    def find_nearest_cluster_id_scaling(self, model: GaussianModel, codebook: torch.Tensor):
+        return self.one_nearst(model.get_scaling.detach(), model.scaling_activation(codebook))
+
+    def produce_clusters(self, model: GaussianModel, init_codebook_dict={}):
+        codebook_dict: Dict[str, torch.Tensor] = {}
+        ids_dict: Dict[str, torch.Tensor] = {}
+        init_codebook_dict = {
+            "features_dc": None,
+            **{f"features_rest_{sh_degree}": None for sh_degree in range(model.max_sh_degree)},
+            "rotation_re": None,
+            "rotation_im": None,
+            "opacity": None,
+            "scaling": None,
+            **init_codebook_dict
+        }
+
+        codebook_dict["features_dc"], ids_dict["features_dc"] = self.produce_clusters_features_dc(model, init_codebook=init_codebook_dict["features_dc"])
+        for sh_degree in range(model.max_sh_degree):
+            codebook_dict[f"features_rest_{sh_degree}"], ids_dict[f"features_rest_{sh_degree}"] = self.produce_clusters_degree_features_rest(
+                model, sh_degree, init_codebook=init_codebook_dict[f"features_rest_{sh_degree}"]
+            )
+        codebook_dict["rotation_re"], ids_dict[f"rotation_re"] = self.produce_clusters_rotation_re(model, init_codebook=init_codebook_dict["rotation_re"])
+        codebook_dict["rotation_im"], ids_dict[f"rotation_im"] = self.produce_clusters_rotation_im(model, init_codebook=init_codebook_dict["rotation_im"])
+        codebook_dict["opacity"], ids_dict[f"opacity"] = self.produce_clusters_opacity(model, init_codebook=init_codebook_dict["opacity"])
+        codebook_dict["scaling"], ids_dict[f"scaling"] = self.produce_clusters_scaling(model, init_codebook=init_codebook_dict["scaling"])
+        return codebook_dict, ids_dict
+
+    def find_nearest_cluster_id(self, model: GaussianModel, codebook_dict={}):
+        ids_dict: Dict[str, torch.Tensor] = {}
+        ids_dict["features_dc"] = self.find_nearest_cluster_id_features_dc(model, codebook=codebook_dict["features_dc"])
+        for sh_degree in range(model.max_sh_degree):
+            ids_dict[f"features_rest_{sh_degree}"] = self.find_nearest_cluster_id_degree_features_rest(
+                model, sh_degree, codebook=codebook_dict[f"features_rest_{sh_degree}"]
+            )
+        ids_dict[f"rotation_re"] = self.find_nearest_cluster_id_rotation_re(model, codebook=codebook_dict["rotation_re"])
+        ids_dict[f"rotation_im"] = self.find_nearest_cluster_id_rotation_im(model, codebook=codebook_dict["rotation_im"])
+        ids_dict[f"opacity"] = self.find_nearest_cluster_id_opacity(model, codebook=codebook_dict["opacity"])
+        ids_dict[f"scaling"] = self.find_nearest_cluster_id_scaling(model, codebook=codebook_dict["scaling"])
+        return ids_dict
+
+    def dequantize(self, model: GaussianModel, ids_dict: Dict[str, torch.Tensor], codebook_dict: Dict[str, torch.Tensor], xyz: torch.Tensor = None, replace=False) -> GaussianModel:
+        opacity = codebook_dict["opacity"][ids_dict["opacity"], ...]
+        scaling = codebook_dict["scaling"][ids_dict["scaling"], ...]
+
+        rotation = torch.cat((
+            codebook_dict["rotation_re"][ids_dict["rotation_re"], ...],
+            codebook_dict["rotation_im"][ids_dict["rotation_im"], ...],
+        ), dim=1)
+
+        features_dc = codebook_dict["features_dc"][ids_dict["features_dc"], ...]
+        features_rest = []
+        for sh_degree in range(model.max_sh_degree):
+            features_rest.append(codebook_dict[f"features_rest_{sh_degree}"][ids_dict[f"features_rest_{sh_degree}"], ...])
+        features_rest = torch.cat(features_rest, dim=2).transpose(1, 2)
+
+        with torch.no_grad():
+            if replace:
+                if xyz is not None:
+                    model._xyz = nn.Parameter(xyz)
+                model._opacity = nn.Parameter(opacity)
+                model._scaling = nn.Parameter(scaling)
+                model._rotation = nn.Parameter(rotation)
+                model._features_dc = nn.Parameter(features_dc)
+                model._features_rest = nn.Parameter(features_rest)
+            else:
+                if xyz is not None:
+                    model._xyz[...] = xyz
+                model._opacity[...] = opacity
+                model._scaling[...] = scaling
+                model._rotation[...] = rotation
+                model._features_dc[...] = features_dc
+                model._features_rest[...] = features_rest
+        return model
+
+    def quantize(self, model: GaussianModel, update_codebook=True) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        if self._codebook_dict == {} or update_codebook:
+            codebook_dict, ids_dict = self.produce_clusters(model, self._codebook_dict)
+            self._codebook_dict = codebook_dict
+        else:
+            codebook_dict = self._codebook_dict
+            ids_dict = self.find_nearest_cluster_id(model, self._codebook_dict)
+        return ids_dict, codebook_dict
+
+    def save_quantized(self, model: GaussianModel, ply_path: str):
+        ids_dict, codebook_dict = self.quantize(model, update_codebook=False)
+        dtype_full = [
+            ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
+            ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'),
+            ('rot_re', self.force_code_dtype or compute_uint_dtype(self.num_clusters_rotation_re)),
+            ('rot_im', self.force_code_dtype or compute_uint_dtype(self.num_clusters_rotation_im)),
+            ('opacity', self.force_code_dtype or compute_uint_dtype(self.num_clusters_opacity)),
+            ('scale', self.force_code_dtype or compute_uint_dtype(self.num_clusters_scaling)),
+            ('f_dc', self.force_code_dtype or compute_uint_dtype(self.num_clusters_features_dc)),
+        ]
+        for sh_degree in range(model.max_sh_degree):
+            force_code_dtype = self.force_code_dtype or compute_uint_dtype(self.num_clusters_features_rest[sh_degree])
+            dtype_full.extend([
+                (f'f_rest_{sh_degree}_0', force_code_dtype),
+                (f'f_rest_{sh_degree}_1', force_code_dtype),
+                (f'f_rest_{sh_degree}_2', force_code_dtype),
+            ])
+        data_full = [
+            *np.array_split(model._xyz.detach().cpu().numpy(), 3, axis=1),
+            *np.array_split(torch.zeros_like(model._xyz).detach().cpu().numpy(), 3, axis=1),
+            ids_dict["rotation_re"].unsqueeze(-1).cpu().numpy(),
+            ids_dict["rotation_im"].unsqueeze(-1).cpu().numpy(),
+            ids_dict["opacity"].unsqueeze(-1).cpu().numpy(),
+            ids_dict["scaling"].unsqueeze(-1).cpu().numpy(),
+            ids_dict["features_dc"].cpu().numpy(),
+        ]
+        for sh_degree in range(model.max_sh_degree):
+            features_rest = ids_dict[f'features_rest_{sh_degree}'].cpu().numpy()
+            data_full.extend(np.array_split(features_rest, 3, axis=1))
+
+        elements = np.rec.fromarrays([data.squeeze(-1) for data in data_full], dtype=dtype_full)
+        el = PlyElement.describe(elements, 'vertex')
+
+        cb = [
+            PlyElement.describe(array2record(codebook_dict["rotation_re"], "rot_re", 1, self.force_codebook_dtype), 'codebook_rot_re'),
+            PlyElement.describe(array2record(codebook_dict["rotation_im"], "rot_im", 3, self.force_codebook_dtype), 'codebook_rot_im'),
+            PlyElement.describe(array2record(codebook_dict["opacity"], "opacity", 1, self.force_codebook_dtype), 'codebook_opacity'),
+            PlyElement.describe(array2record(codebook_dict["scaling"], "scaling", 3, self.force_codebook_dtype), 'codebook_scaling'),
+            PlyElement.describe(array2record(codebook_dict["features_dc"], "f_dc", 3, self.force_codebook_dtype), 'codebook_f_dc'),
+        ]
+        for sh_degree in range(model.max_sh_degree):
+            features_rest = codebook_dict[f'features_rest_{sh_degree}']
+            n_channels = (sh_degree + 2) ** 2 - (sh_degree + 1) ** 2
+            cb.append(PlyElement.describe(array2record(features_rest, f'f_rest_{sh_degree}', n_channels, self.force_codebook_dtype), f'codebook_f_rest_{sh_degree}'))
+
+        PlyData([el, *cb]).write(ply_path)
+
+    def load_quantized(self, model: GaussianModel, ply_path: str) -> GaussianModel:
+        plydata = PlyData.read(ply_path)
+
+        ids_dict = {}
+        elements = plydata['vertex']
+        kwargs = dict(dtype=torch.long, device=model._xyz.device)
+        ids_dict["rotation_re"] = torch.tensor(elements["rot_re"].copy(), **kwargs)
+        ids_dict["rotation_im"] = torch.tensor(elements["rot_im"].copy(), **kwargs)
+        ids_dict["opacity"] = torch.tensor(elements["opacity"].copy(), **kwargs)
+        ids_dict["scaling"] = torch.tensor(elements["scale"].copy(), **kwargs)
+        ids_dict["features_dc"] = torch.tensor(elements["f_dc"].copy(), **kwargs).unsqueeze(-1)
+        for sh_degree in range(model.max_sh_degree):
+            ids_dict[f'features_rest_{sh_degree}'] = torch.tensor(np.stack([elements[f'f_rest_{sh_degree}_{ch}'] for ch in range(3)], axis=1), **kwargs)
+
+        codebook_dict = {}
+        kwargs = dict(dtype=torch.float32, device=model._xyz.device)
+        codebook_dict["rotation_re"] = torch.tensor(plydata["codebook_rot_re"]["rot_re"], **kwargs).unsqueeze(-1)
+        codebook_dict["rotation_im"] = torch.tensor(np.stack([plydata["codebook_rot_im"][f'rot_im_{ch}'] for ch in range(3)], axis=1), **kwargs)
+        codebook_dict["opacity"] = torch.tensor(plydata["codebook_opacity"]["opacity"], **kwargs).unsqueeze(-1)
+        codebook_dict["scaling"] = torch.tensor(np.stack([plydata["codebook_scaling"][f'scaling_{ch}'] for ch in range(3)], axis=1), **kwargs)
+        codebook_dict["features_dc"] = torch.tensor(np.stack([plydata["codebook_f_dc"][f'f_dc_{ch}'] for ch in range(3)], axis=1), **kwargs)
+        for sh_degree in range(model.max_sh_degree):
+            n_channels = (sh_degree + 2) ** 2 - (sh_degree + 1) ** 2
+            codebook_dict[f'features_rest_{sh_degree}'] = torch.tensor(np.stack([plydata[f"codebook_f_rest_{sh_degree}"][f'f_rest_{sh_degree}_{ch}'] for ch in range(n_channels)], axis=1), **kwargs)
+
+        self._codebook_dict = codebook_dict
+
+        xyz = torch.stack([
+            torch.tensor(elements["x"].copy(), **kwargs),
+            torch.tensor(elements["y"].copy(), **kwargs),
+            torch.tensor(elements["z"].copy(), **kwargs),
+        ], dim=1)
+        return self.dequantize(model, ids_dict, codebook_dict, xyz=xyz, replace=True)

reduced_3dgs/quantization/wrapper.py

@@ -0,0 +1,67 @@
+from gaussian_splatting import GaussianModel
+from gaussian_splatting.trainer import AbstractTrainer, Trainer
+from .abc import QuantizeTrainerWrapper
+from .exclude_zeros import ExcludeZeroSHQuantizer
+
+
+def VectorQuantizeTrainerWrapper(
+    base_trainer: AbstractTrainer,
+        num_clusters=256,
+        num_clusters_rotation_re=None,
+        num_clusters_rotation_im=None,
+        num_clusters_opacity=None,
+        num_clusters_scaling=None,
+        num_clusters_features_dc=None,
+        num_clusters_features_rest=[],
+        quantize_from_iter=5000,
+        quantize_until_iter=30000,
+        quantize_interval=1000,
+        treat_as_zero=1e-8,
+):
+    return QuantizeTrainerWrapper(
+        base_trainer, ExcludeZeroSHQuantizer(
+            num_clusters=num_clusters,
+            num_clusters_rotation_re=num_clusters_rotation_re,
+            num_clusters_rotation_im=num_clusters_rotation_im,
+            num_clusters_opacity=num_clusters_opacity,
+            num_clusters_scaling=num_clusters_scaling,
+            num_clusters_features_dc=num_clusters_features_dc,
+            num_clusters_features_rest=num_clusters_features_rest,
+            treat_as_zero=treat_as_zero,
+        ),
+        quantize_from_iter=quantize_from_iter,
+        quantize_until_iter=quantize_until_iter,
+        quantize_interval=quantize_interval,
+
+    )
+
+
+def VectorQuantizeTrainer(
+    model: GaussianModel,
+    scene_extent: float,
+        num_clusters=256,
+        num_clusters_rotation_re=None,
+        num_clusters_rotation_im=None,
+        num_clusters_opacity=None,
+        num_clusters_scaling=None,
+        num_clusters_features_dc=None,
+        num_clusters_features_rest=[],
+        quantize_from_iter=5000,
+        quantize_until_iter=30000,
+        quantize_interval=1000,
+        treat_as_zero=1e-8,
+        *args, **kwargs):
+    return VectorQuantizeTrainerWrapper(
+        Trainer(model, scene_extent, *args, **kwargs),
+        num_clusters=num_clusters,
+        num_clusters_rotation_re=num_clusters_rotation_re,
+        num_clusters_rotation_im=num_clusters_rotation_im,
+        num_clusters_opacity=num_clusters_opacity,
+        num_clusters_scaling=num_clusters_scaling,
+        num_clusters_features_dc=num_clusters_features_dc,
+        num_clusters_features_rest=num_clusters_features_rest,
+        treat_as_zero=treat_as_zero,
+        quantize_from_iter=quantize_from_iter,
+        quantize_until_iter=quantize_until_iter,
+        quantize_interval=quantize_interval,
+    )

reduced_3dgs/quantize.py

@@ -0,0 +1,49 @@
+import os
+import shutil
+from gaussian_splatting import GaussianModel
+from reduced_3dgs.quantization import ExcludeZeroSHQuantizer as VectorQuantizer
+
+
+def copy_not_exists(source, destination):
+    if os.path.exists(destination):
+        if os.path.samefile(source, destination):
+            return
+        os.remove(destination)
+    shutil.copy(source, destination)
+
+
+def quantize(source, destination, iteration, sh_degree, device, **kwargs):
+    copy_not_exists(os.path.join(source, "cfg_args"), os.path.join(destination, "cfg_args"))
+    copy_not_exists(os.path.join(source, "cameras.json"), os.path.join(destination, "cameras.json"))
+
+    input = os.path.join(source, "point_cloud", "iteration_" + str(iteration), "point_cloud.ply")
+    output = os.path.join(destination, "point_cloud", "iteration_" + str(iteration), "point_cloud_quantized.ply")
+    os.makedirs(os.path.join(destination, "point_cloud", "iteration_" + str(iteration)), exist_ok=True)
+    gaussians = GaussianModel(sh_degree).to(device)
+    gaussians.load_ply(input)
+    quantizer = VectorQuantizer(**kwargs)
+    quantizer.save_quantized(gaussians, output)
+    ids_dict, codebook_dict = quantizer.quantize(gaussians, update_codebook=False)
+    gaussians = quantizer.dequantize(gaussians, ids_dict, codebook_dict)
+    quantizer.load_quantized(gaussians, output)
+    output = os.path.join(destination, "point_cloud", "iteration_" + str(iteration), "point_cloud.ply")
+    gaussians.save_ply(output)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-s", "--source", required=True, type=str)
+    parser.add_argument("-d", "--destination", required=True, type=str)
+    parser.add_argument("-i", "--iteration", default=30000, type=int)
+    parser.add_argument("--sh_degree", type=int, default=3)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--num_clusters", type=int, default=256)
+    parser.add_argument("--num_clusters_rotation_re", type=int, default=None)
+    parser.add_argument("--num_clusters_rotation_im", type=int, default=None)
+    parser.add_argument("--num_clusters_opacity", type=int, default=None)
+    parser.add_argument("--num_clusters_scaling", type=int, default=None)
+    parser.add_argument("--num_clusters_features_dc", type=int, default=None)
+    parser.add_argument("--num_clusters_features_rest", nargs="+", type=int, default=[])
+    args = parser.parse_args()
+    quantize(**vars(args))

reduced_3dgs/shculling/__init__.py

@@ -0,0 +1,2 @@
+from .gaussian_model import VariableSHGaussianModel
+from .trainer import SHCuller, SHCullingTrainerWrapper, BaseSHCullingTrainer, SHCullingTrainer