reduced-3dgs 1.10.1__cp312-cp312-win_amd64.whl → 1.10.3__cp312-cp312-win_amd64.whl

reduced_3dgs/prepare.py

@@ -0,0 +1,105 @@
+from gaussian_splatting import GaussianModel
+from gaussian_splatting.dataset import CameraDataset
+from gaussian_splatting.dataset.colmap import colmap_init
+from gaussian_splatting.trainer import AbstractTrainer
+from gaussian_splatting.trainer.extensions import ScaleRegularizeTrainerWrapper
+from reduced_3dgs.quantization import VectorQuantizeTrainerWrapper
+from reduced_3dgs.shculling import VariableSHGaussianModel, SHCullingTrainer
+from reduced_3dgs.pruning import PruningTrainer
+from reduced_3dgs.combinations import PrunerInDensifyTrainer, SHCullingDensificationTrainer, SHCullingPruningTrainer, SHCullingPrunerInDensifyTrainer
+from reduced_3dgs.combinations import CameraTrainableVariableSHGaussianModel, CameraSHCullingTrainer, CameraPruningTrainer
+from reduced_3dgs.combinations import CameraPrunerInDensifyTrainer, CameraSHCullingDensifyTrainer, CameraSHCullingPruningTrainer, CameraSHCullingPruningDensifyTrainer
+
+
+def prepare_gaussians(sh_degree: int, source: str, device: str, trainable_camera: bool = False, load_ply: str = None) -> GaussianModel:
+    if trainable_camera:
+        gaussians = CameraTrainableVariableSHGaussianModel(sh_degree).to(device)
+        gaussians.load_ply(load_ply) if load_ply else colmap_init(gaussians, source)
+    else:
+        gaussians = VariableSHGaussianModel(sh_degree).to(device)
+        gaussians.load_ply(load_ply) if load_ply else colmap_init(gaussians, source)
+    return gaussians
+
+
+modes = {
+    "shculling": SHCullingTrainer,
+    "pruning": PruningTrainer,
+    "densify-pruning": PrunerInDensifyTrainer,
+    "densify-shculling": SHCullingDensificationTrainer,
+    "prune-shculling": SHCullingPruningTrainer,
+    "densify-prune-shculling": SHCullingPrunerInDensifyTrainer,
+    "camera-shculling": CameraSHCullingTrainer,
+    "camera-pruning": CameraPruningTrainer,
+    "camera-densify-pruning": CameraPrunerInDensifyTrainer,
+    "camera-densify-shculling": CameraSHCullingDensifyTrainer,
+    "camera-prune-shculling": CameraSHCullingPruningTrainer,
+    "camera-densify-prune-shculling": CameraSHCullingPruningDensifyTrainer,
+}
+
+
+def prepare_quantizer(
+        gaussians: GaussianModel,
+        scene_extent: float,
+        dataset: CameraDataset,
+        base_constructor,
+        load_quantized: str = None,
+
+        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,
+        **configs):
+    trainer = VectorQuantizeTrainerWrapper(
+        base_constructor(
+            gaussians,
+            scene_extent=scene_extent,
+            dataset=dataset,
+            **configs
+        ),
+
+        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,
+
+        quantize_from_iter=quantize_from_iter,
+        quantize_until_iter=quantize_until_iter,
+        quantize_interval=quantize_interval,
+    )
+    if load_quantized:
+        trainer.quantizer.load_quantized(load_quantized)
+    return trainer, trainer.quantizer
+
+
+def prepare_trainer(gaussians: GaussianModel, dataset: CameraDataset, mode: str, with_scale_reg=False, quantize: bool = False, load_quantized: str = None, configs={}) -> AbstractTrainer:
+    constructor = modes[mode]
+    if with_scale_reg:
+        constructor = lambda *args, **kwargs: ScaleRegularizeTrainerWrapper(modes[mode], *args, **kwargs)
+    if quantize:
+        trainer, quantizer = prepare_quantizer(
+            gaussians,
+            scene_extent=dataset.scene_extent(),
+            dataset=dataset,
+            base_constructor=modes[mode],
+            load_quantized=load_quantized,
+            **configs
+        )
+    else:
+        trainer = constructor(
+            gaussians,
+            scene_extent=dataset.scene_extent(),
+            dataset=dataset,
+            **configs
+        )
+        quantizer = None
+    return trainer, quantizer

reduced_3dgs/quantization/quantizer.py

@@ -206,8 +206,7 @@ class VectorQuantizer(AbstractQuantizer):
             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)
+    def ply_dtype(self, max_sh_degree: int):
         dtype_full = [
             ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
             ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4'),
@@ -217,13 +216,16 @@ class VectorQuantizer(AbstractQuantizer):
             ('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):
+        for sh_degree in range(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),
             ])
+        return dtype_full
+
+    def ply_data(self, model: GaussianModel, ids_dict: Dict[str, torch.Tensor]):
         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),
@@ -236,6 +238,12 @@ class VectorQuantizer(AbstractQuantizer):
         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))
+        return data_full
+
+    def save_quantized(self, model: GaussianModel, ply_path: str):
+        ids_dict, codebook_dict = self.quantize(model, update_codebook=False)
+        dtype_full = self.ply_dtype(model.max_sh_degree)
+        data_full = self.ply_data(model, ids_dict)
 
         elements = np.rec.fromarrays([data.squeeze(-1) for data in data_full], dtype=dtype_full)
         el = PlyElement.describe(elements, 'vertex')

reduced_3dgs/quantize.py

@@ -13,6 +13,7 @@ def copy_not_exists(source, destination):
 
 
 def quantize(source, destination, iteration, sh_degree, device, **kwargs):
+    os.makedirs(destination, exist_ok=True)
     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"))
 

reduced_3dgs/train.py

@@ -1,129 +1,22 @@
 import os
 import random
 import shutil
-from typing import List, Tuple
+from typing import List
 import torch
 from tqdm import tqdm
 from gaussian_splatting import GaussianModel
 from gaussian_splatting.dataset import CameraDataset
 from gaussian_splatting.utils import psnr
-from gaussian_splatting.dataset.colmap import colmap_init
 from gaussian_splatting.trainer import AbstractTrainer
-from gaussian_splatting.trainer.extensions import ScaleRegularizeTrainerWrapper
-from gaussian_splatting.train import prepare_dataset, save_cfg_args
-from reduced_3dgs.quantization import AbstractQuantizer, VectorQuantizeTrainerWrapper
-from reduced_3dgs.shculling import VariableSHGaussianModel, SHCullingTrainer
-from reduced_3dgs.pruning import PruningTrainer
-from reduced_3dgs.combinations import PrunerInDensifyTrainer, SHCullingDensificationTrainer, SHCullingPruningTrainer, SHCullingPrunerInDensifyTrainer
-from reduced_3dgs.combinations import CameraTrainableVariableSHGaussianModel, CameraSHCullingTrainer, CameraPruningTrainer
-from reduced_3dgs.combinations import CameraPrunerInDensifyTrainer, CameraSHCullingDensifyTrainer, CameraSHCullingPruningTrainer, CameraSHCullingPruningDensifyTrainer
+from gaussian_splatting.prepare import prepare_dataset
+from gaussian_splatting.train import save_cfg_args
+from reduced_3dgs.quantization import AbstractQuantizer
+from reduced_3dgs.prepare import modes, prepare_gaussians, prepare_trainer
 
 
-basemodes = {
-    "shculling": SHCullingTrainer,
-    "pruning": PruningTrainer,
-    "densify-pruning": PrunerInDensifyTrainer,
-    "densify-shculling": SHCullingDensificationTrainer,
-    "prune-shculling": SHCullingPruningTrainer,
-    "densify-prune-shculling": SHCullingPrunerInDensifyTrainer,
-}
-cameramodes = {
-    "camera-shculling": CameraSHCullingTrainer,
-    "camera-pruning": CameraPruningTrainer,
-    "camera-densify-pruning": CameraPrunerInDensifyTrainer,
-    "camera-densify-shculling": CameraSHCullingDensifyTrainer,
-    "camera-prune-shculling": CameraSHCullingPruningTrainer,
-    "camera-densify-prune-shculling": CameraSHCullingPruningDensifyTrainer,
-}
-
-
-def prepare_quantizer(
-        gaussians: GaussianModel,
-        scene_extent: float,
-        dataset: CameraDataset,
-        base_constructor,
-        load_quantized: str = None,
-
-        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,
-        **configs):
-    trainer = VectorQuantizeTrainerWrapper(
-        base_constructor(
-            gaussians,
-            scene_extent=scene_extent,
-            dataset=dataset,
-            **configs
-        ),
-
-        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,
-
-        quantize_from_iter=quantize_from_iter,
-        quantize_until_iter=quantize_until_iter,
-        quantize_interval=quantize_interval,
-    )
-    if load_quantized:
-        trainer.quantizer.load_quantized(load_quantized)
-    return trainer, trainer.quantizer
-
-
-def prepare_gaussians(sh_degree: int, source: str, device: str, trainable_camera: bool = False, load_ply: str = None) -> Tuple[CameraDataset, GaussianModel, AbstractTrainer]:
-    if trainable_camera:
-        gaussians = CameraTrainableVariableSHGaussianModel(sh_degree).to(device)
-        gaussians.load_ply(load_ply) if load_ply else colmap_init(gaussians, source)
-    else:
-        gaussians = VariableSHGaussianModel(sh_degree).to(device)
-        gaussians.load_ply(load_ply) if load_ply else colmap_init(gaussians, source)
-    return gaussians
-
-
-def prepare_trainer(gaussians: GaussianModel, dataset: CameraDataset, mode: str, with_scale_reg=False, quantize: bool = False, load_quantized: str = None, configs={}) -> AbstractTrainer:
-    if mode in basemodes:
-        modes = basemodes
-    elif mode in cameramodes:
-        modes = cameramodes
-    else:
-        raise ValueError(f"Unknown mode: {mode}")
-    constructor = modes[mode]
-    if with_scale_reg:
-        constructor = lambda *args, **kwargs: ScaleRegularizeTrainerWrapper(modes[mode], *args, **kwargs)
-    if quantize:
-        trainer, quantizer = prepare_quantizer(
-            gaussians,
-            scene_extent=dataset.scene_extent(),
-            dataset=dataset,
-            base_constructor=modes[mode],
-            load_quantized=load_quantized,
-            **configs
-        )
-    else:
-        trainer = constructor(
-            gaussians,
-            scene_extent=dataset.scene_extent(),
-            dataset=dataset,
-            **configs
-        )
-        quantizer = None
-    return trainer, quantizer
-
-
-def prepare_training(sh_degree: int, source: str, device: str, mode: str, load_ply: str = None, load_camera: str = None, load_depth=False, with_scale_reg=False, quantize: bool = False, load_quantized: str = None, configs={}) -> Tuple[CameraDataset, GaussianModel, AbstractTrainer]:
-    dataset = prepare_dataset(source=source, device=device, trainable_camera=mode in cameramodes, load_camera=load_camera, load_depth=load_depth)
-    gaussians = prepare_gaussians(sh_degree=sh_degree, source=source, device=device, trainable_camera=mode in cameramodes, load_ply=load_ply)
+def prepare_training(sh_degree: int, source: str, device: str, mode: str, trainable_camera: bool = False, load_ply: str = None, load_camera: str = None, load_depth=False, with_scale_reg=False, quantize: bool = False, load_quantized: str = None, configs={}):
+    dataset = prepare_dataset(source=source, device=device, trainable_camera=trainable_camera, load_camera=load_camera, load_depth=load_depth)
+    gaussians = prepare_gaussians(sh_degree=sh_degree, source=source, device=device, trainable_camera=trainable_camera, load_ply=load_ply)
     trainer, quantizer = prepare_trainer(gaussians=gaussians, dataset=dataset, mode=mode, with_scale_reg=with_scale_reg, quantize=quantize, load_quantized=load_quantized, configs=configs)
     return dataset, gaussians, trainer, quantizer
 
@@ -166,7 +59,7 @@ def training(dataset: CameraDataset, gaussians: GaussianModel, trainer: Abstract
 
 
 if __name__ == "__main__":
-    from argparse import ArgumentParser, Namespace
+    from argparse import ArgumentParser
     parser = ArgumentParser()
     parser.add_argument("--sh_degree", default=3, type=int)
     parser.add_argument("-s", "--source", required=True, type=str)
@@ -178,7 +71,7 @@ if __name__ == "__main__":
     parser.add_argument("--no_depth_data", action='store_true')
     parser.add_argument("--with_scale_reg", action="store_true")
     parser.add_argument("--load_quantized", default=None, type=str)
-    parser.add_argument("--mode", choices=list(basemodes.keys()) + list(cameramodes.keys()), default="densify-prune-shculling")
+    parser.add_argument("--mode", choices=list(modes), default="densify-prune-shculling")
     parser.add_argument("--save_iterations", nargs="+", type=int, default=[7000, 30000])
     parser.add_argument("--device", default="cuda", type=str)
     parser.add_argument("--empty_cache_every_step", action='store_true')
@@ -189,7 +82,7 @@ if __name__ == "__main__":
 
     configs = {o.split("=", 1)[0]: eval(o.split("=", 1)[1]) for o in args.option}
     dataset, gaussians, trainer, quantizer = prepare_training(
-        sh_degree=args.sh_degree, source=args.source, device=args.device, mode=args.mode,
+        sh_degree=args.sh_degree, source=args.source, device=args.device, mode=args.mode, trainable_camera="camera" in args.mode,
         load_ply=args.load_ply, load_camera=args.load_camera, load_depth=not args.no_depth_data, with_scale_reg=args.with_scale_reg,
         quantize=args.quantize, load_quantized=args.load_quantized, configs=configs)
     dataset.save_cameras(os.path.join(args.destination, "cameras.json"))

{reduced_3dgs-1.10.1.dist-info → reduced_3dgs-1.10.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: reduced_3dgs
-Version: 1.10.1
+Version: 1.10.3
 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.10.1.dist-info → reduced_3dgs-1.10.3.dist-info}/RECORD

@@ -1,13 +1,14 @@
 reduced_3dgs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 reduced_3dgs/combinations.py,sha256=k4ErxpAscCqJMtVXZ29KGQjw2VoZMV0W3V4u3nj0e-Y,9401
-reduced_3dgs/quantize.py,sha256=Y44qHyFdOIqke7NoeqXmyKloS43j-al74ZiNsuZZHbM,2527
-reduced_3dgs/train.py,sha256=kXgUZqSGc09NNmGNCoSe3Dv-VMBuQ6XjVNJhjmm0g3E,10082
-reduced_3dgs/diff_gaussian_rasterization/_C.cp312-win_amd64.pyd,sha256=abu4zXWAobWzWcuvBTv5kQqqe6TJ4RDm817Mjdpli5w,1640448
+reduced_3dgs/prepare.py,sha256=MFUUckRHKfgcva4ZOBxfPFyE95N-OlCQLplpmEPuzOk,4440
+reduced_3dgs/quantize.py,sha256=BVqBb2tQgiP3hap5-OByD8VELtJJGfEeFzaVFyzCJZU,2572
+reduced_3dgs/train.py,sha256=yRQPQv-hTOBQN-tqGHYs2aIZ0AbWq158CXpthYN2cfw,5666
+reduced_3dgs/diff_gaussian_rasterization/_C.cp312-win_amd64.pyd,sha256=1n2kPV5fkgfh94MJXmKd1-f-aQo2FdqXPr8g8jMsngQ,1640448
 reduced_3dgs/diff_gaussian_rasterization/__init__.py,sha256=oV6JjTc-50MscX4XHeIWSgLr3l8Y25knBIs-0gRbJr4,7932
 reduced_3dgs/importance/__init__.py,sha256=neJsbY5cLikEGBQGdR4MjwCQ5VWVikT1357DwL0EtWU,289
 reduced_3dgs/importance/combinations.py,sha256=eAdykeTdvRGCHxskjILQnZVaqQVvwC-0wMxdgYMeeDs,2922
 reduced_3dgs/importance/trainer.py,sha256=Sj4ORvoYtFT7z3hifzFZDfhFyqumHraXyk3vMVtk0AU,12661
-reduced_3dgs/importance/diff_gaussian_rasterization/_C.cp312-win_amd64.pyd,sha256=5ZxCa8jgG1bC4Ub77UrmNv-iLxvxTyxokQQVkqTNHZc,1320448
+reduced_3dgs/importance/diff_gaussian_rasterization/_C.cp312-win_amd64.pyd,sha256=2TshQzpg0ERxwfB6M9qQfp14KW4QVsxpmEtmHoT2nFw,1320448
 reduced_3dgs/importance/diff_gaussian_rasterization/__init__.py,sha256=Tix8auyXBb_QFQtXrV3sLE9kdnl5zgHH0BbqcFzDp84,12850
 reduced_3dgs/pruning/__init__.py,sha256=E_YxJ9cDV_B6EJbYUBEcuRYMIht_C72rI1VJUXFCLpM,201
 reduced_3dgs/pruning/combinations.py,sha256=QhXt2C7pTXhwzp9hPL9dVdiQzz0cUQpm5qljqytPEsM,2345
@@ -15,14 +16,14 @@ reduced_3dgs/pruning/trainer.py,sha256=JJml-uYfDfUpbsjRNZbIvnUYYslVgFXkhejbkYSo0
 reduced_3dgs/quantization/__init__.py,sha256=1z1xMn3yj9u7cR9JizGrI3WSyIES_Tqq6oDquvglSeo,225
 reduced_3dgs/quantization/abc.py,sha256=rsi8HFRwQCltWTYiJ3BpygtQDT7hK6J01jKMOboOY8w,1910
 reduced_3dgs/quantization/exclude_zeros.py,sha256=fKSgjHous4OpdI6mQi9z23if9jnbB79w2jChpxkCJWw,2381
-reduced_3dgs/quantization/quantizer.py,sha256=VL96ZTY7G8jUt-fLBJ3gw1MQV46NkHOBJyfUPsmDQxE,16712
+reduced_3dgs/quantization/quantizer.py,sha256=HZFuYs1Uxq3aXzn9Ow1LSDxF4myh_vfgvovVDJoc2BM,17001
 reduced_3dgs/quantization/wrapper.py,sha256=cyXqfJgo9b3fS7DYXxOk5LmQudvrEhweOebFsjRnXiQ,2549
 reduced_3dgs/shculling/__init__.py,sha256=nP2BejDCUdCmJNRbg0hfhHREO6jyZXwIcRiw6ttVgqo,149
 reduced_3dgs/shculling/gaussian_model.py,sha256=f8QWaL09vaV9Tcf6Dngjg_Fmk1wTQPAjWhuhI_N02Y8,2877
 reduced_3dgs/shculling/trainer.py,sha256=9hwR77djhZpyf-URhwKHjnLbe0ZAOS-DIw58RzkcHXQ,6369
-reduced_3dgs/simple_knn/_C.cp312-win_amd64.pyd,sha256=RQosX1BwSnkCObjwlwRsw2fnDDR7IwHyVOp4cx8IX_Q,1267712
-reduced_3dgs-1.10.1.dist-info/licenses/LICENSE.md,sha256=LQ4_LAqlncGkg_mQy5ykMAFtQDSPB0eKmIEtBut0yjw,4916
-reduced_3dgs-1.10.1.dist-info/METADATA,sha256=6cxgTTpPqokRIa7VNyuXEJ2wN5iyClcSwbY451onMrA,13015
-reduced_3dgs-1.10.1.dist-info/WHEEL,sha256=8UP9x9puWI0P1V_d7K2oMTBqfeLNm21CTzZ_Ptr0NXU,101
-reduced_3dgs-1.10.1.dist-info/top_level.txt,sha256=PpU5aT3-baSCdqCtTaZknoB32H93UeKCkYDkRCCZMEI,13
-reduced_3dgs-1.10.1.dist-info/RECORD,,
+reduced_3dgs/simple_knn/_C.cp312-win_amd64.pyd,sha256=5ZzPxMlTNRnsLs2p9EFSr93-MI9XQEUFgFo8SRhazxo,1267712
+reduced_3dgs-1.10.3.dist-info/licenses/LICENSE.md,sha256=LQ4_LAqlncGkg_mQy5ykMAFtQDSPB0eKmIEtBut0yjw,4916
+reduced_3dgs-1.10.3.dist-info/METADATA,sha256=DGxZflgGcJRjRyeqz0CTkgSUeJpGGAgKDkOJ0mG5Dk4,13015
+reduced_3dgs-1.10.3.dist-info/WHEEL,sha256=8UP9x9puWI0P1V_d7K2oMTBqfeLNm21CTzZ_Ptr0NXU,101
+reduced_3dgs-1.10.3.dist-info/top_level.txt,sha256=PpU5aT3-baSCdqCtTaZknoB32H93UeKCkYDkRCCZMEI,13
+reduced_3dgs-1.10.3.dist-info/RECORD,,