bugfix: bugs coming from branch merging

Author: JasonLSC

README.md

@@ -25,7 +25,8 @@ cd examples
 pip install -r requirements.txt
 # download mipnerf_360 benchmark data
 python datasets/download_dataset.py
-# place other dataset under 'data' folder
+# or place other dataset under 'data' folder
+ln -s data/tandt /xxxx/Dataset/tandt
 ```
 
 We also use third-party library, 'python-fpnge', to accelerate image saving operations during the experiment for now. 

gsplat/compression/entropy_coding_compression.py

@@ -13,12 +13,7 @@
 from gsplat.compression.sort import sort_splats
 from gsplat.utils import inverse_log_transform, log_transform
 
-try:
-    import constriction
-except:
-    raise ImportError(
-        "Please install constriction with 'pip install constriction' to use ANS"
-    )
+
 
 
 @dataclass
@@ -262,6 +257,13 @@ def _get_likelihood(symbols: np.array, bitwidth: int=8) -> np.array:
     pass
 
 def _categorical_ans_encode(symbols: np.array, probabilities: np.array, save_path:str):
+    try:
+        import constriction
+    except:
+        raise ImportError(
+            "Please install constriction with 'pip install constriction' to use ANS"
+    )
+
     num_symbols = symbols.shape[-1]
 
     message_list = []
@@ -364,7 +366,13 @@ def _decompress_factorized_ans(compress_dir: str, param_name: str, meta: Dict[st
     Returns:
         Tensor: parameters
     """
-    import imageio.v2 as imageio
+    try:
+        import constriction
+    except:
+        raise ImportError(
+            "Please install constriction with 'pip install constriction' to use ANS"
+        )
+
     if param_name == "sh0":
         import pdb; pdb.set_trace()
     if not np.all(meta["shape"]):
@@ -413,6 +421,12 @@ def _compress_gaussian_ans(
     Returns:
         Dict[str, Any]: metadata
     """
+    try:
+        import constriction
+    except:
+        raise ImportError(
+            "Please install constriction with 'pip install constriction' to use ANS"
+        )
 
     mins = torch.amin(params, dim=0)
     maxs = torch.amax(params, dim=0)
@@ -683,7 +697,7 @@ def _compress_kmeans(
             "dtype": str(params.dtype).split(".")[1],
         }
         return meta
-    # import pdb; pdb.set_trace()
+    
     kmeans = KMeans(n_clusters=n_clusters, distance="manhattan", verbose=verbose)
     x = params.reshape(params.shape[0], -1).permute(1, 0).contiguous()
     labels = kmeans.fit(x)

gsplat/compression_simulation/gaussian_distribution_model.py

@@ -8,7 +8,7 @@
     import _gridencoder as _backend
 except:
     raise ImportError(
-            "Please install gridencoder with 'pip install gscodec/gridencoder' to use hash encoding"
+            "Please install gridencoder with 'pip install third_party/gridencoder' to use hash encoding"
         )
 
 anchor_round_digits = 16

gsplat/compression_simulation/simulation.py

@@ -502,43 +502,285 @@ def forward(ctx, input):
     @staticmethod
     def backward(ctx, grad_output):
         return grad_output
+
+
+
+class STGCompressionSimulation:
+    """
+    """
+    def __init__(self, quantization_sim_type: Optional[Literal["round", "noise", "vq"]] = None,
+                 entropy_model_enable: bool = False,
+                 entropy_steps: Dict[str, int] = None, 
+                 device: device = None, 
+                 ada_mask_opt: bool = False,
+                 ada_mask_step: int = 10_000,
+                 **kwargs) -> None:
+        self.quantization_sim_type = quantization_sim_type
+
+        self.entropy_model_enable = entropy_model_enable
+        self.entropy_steps = entropy_steps
+        self.device = device
+
+        # simulation_option: dict to specify which properties should be involved in the compression simulation.
+        # Once option is set to True, it must have corresponding simulate_fn
+        self.simulation_option = {
+            "means": False,
+            "scales": True,
+            "quats": True,
+            "opacities": True,
+            "trbf_center": False,
+            "trbf_scale": False,
+            "motion": False, # [N, 9]
+            "omega": False, # [N, 4]
+            "colors": True,
+            "features_dir": True,
+            "features_time": True,
+        }
+
+        self.shN_qat = False
+        self.shN_ada_mask_opt = ada_mask_opt
+        self.shN_ada_mask_step = ada_mask_step
+
+        # configs for "differentiable quantization"
+        self.q_bitwidth = {
+            "means": None,
+            "scales": 8,
+            "quats": 8,
+            "opacities": 8,
+            "trbf_center": None,
+            "trbf_scale": None,
+            "motion": None, # [N, 9]
+            "omega": None, # [N, 4]
+            "colors": 8,
+            "features_dir": 8,
+            "features_time": 8,
+        }
+
+        self.bds = {
+            "means": None,
+            "scales": [-10, 2],
+            "quats": [-1, 1],
+            "opacities": [-7, 7],
+            "trbf_center": None,
+            "trbf_scale": None,
+            "motion": None, # [N, 9]
+            "omega": None, # [N, 4]
+            "colors": [-7.5, 7.5],
+            "features_dir": [-10, 10],
+            "features_time": [-10, 10],
+        }
+
+        # configs for "entropy constraint"
+        self.entropy_model_option = {
+            "means": False,
+            "scales": True,
+            "quats": True,
+            "opacities": False,
+            "colors": True,
+            "features_dir": True,
+            "features_time": True
+            # "shN": False
+        }
+
+        if self.entropy_model_enable:
+            self.entropy_models = {
+                "means": None,
+                "scales": Entropy_factorized_optimized_refactor(channel=3).to(self.device),
+                # "scales": None,
+                "quats": Entropy_factorized_optimized_refactor(channel=4).to(self.device),
+                "opacities": None,
+                "colors": Entropy_factorized_optimized_refactor(channel=3, filters=(3, 3)).to(self.device),
+                "features_dir": Entropy_factorized_optimized_refactor(channel=3, filters=(3, 3)).to(self.device),
+                "features_time": Entropy_factorized_optimized_refactor(channel=3, filters=(3, 3)).to(self.device),
+            }
+
+            self.entropy_model_optimizers = {}
+            for k, v in self.entropy_models.items():
+                if isinstance(v, Entropy_factorized) or isinstance(v, Entropy_factorized_optimized) or isinstance(v, Entropy_factorized_optimized_refactor):
+                    v_opt = torch.optim.Adam(
+                        [{"params": p, "lr": 1e-4, "name": n} for n, p in v.named_parameters()]
+                    )
+                    # v_opt = torch.optim.SGD(
+                    #     [{"params": p, "lr": 1e-4, "name": n} for n, p in v.named_parameters()]
+                    # )
+                else:
+                    v_opt = None
+                self.entropy_model_optimizers.update({k: v_opt})
+
+        # configs for "adaptive mask"
+        if self.shN_ada_mask_opt:
+            from .ada_mask import AnnealingMask
+            cap_max = kwargs.get("cap_max", 1_000_000)
+            self.shN_ada_mask = AnnealingMask(input_shape=[cap_max, 1, 1], 
+                                              device=device,
+                                              annealing_start_iter=ada_mask_step)
+            
+            self.shN_ada_mask_optimizer = torch.optim.Adam([
+                {'params': self.shN_ada_mask.parameters(), 'lr': 0.01}
+            ])
+
+    def _get_simulate_fn(self, param_name: str) -> Callable:
+        simulate_fn_map = {
+            "means": self.simulate_compression_means,
+            "scales": self.simulate_compression_scales,
+            "quats": self.simulate_compression_quats,
+            "opacities": self.simulate_compression_opacities,
+            # "trbf_center": self.simulate_compression_trbf_center,
+            # "trbf_scale": self.simulate_compression_trbf_scale,
+            # "motion": self.simulate_compression_motion,
+            # "omega": self.simulate_compression_omega,
+            "colors": self.simulate_compression_colors,
+            "features_dir": self.simulate_compression_features_dir,
+            "features_time": self.simulate_compression_features_time
+        }
+        if param_name in simulate_fn_map:
+            return simulate_fn_map[param_name]
+        else:
+            return torch.nn.Identity()
+
+    def simulate_compression(self, splats: Dict[str, Tensor], step: int) -> Dict[str, Tensor]:
+        """
+        """
+        # Create empty dicts for output, including fake quantized values and (optional) estimated bits
+        new_splats = {}
+        esti_bits_dict = {}
+
+        # # Randomly sample approximately 5% of the points rather than all points for speedup.
+        # choose_idx = torch.rand_like(splats["means"][:, 0], device=self.device) <= 1
+        choose_idx = None
+
+        for param_name in splats.keys():
+            # Check which params need to be simulate 
+            if self.simulation_option[param_name]:
+                simulate_fn = self._get_simulate_fn(param_name)
+                new_splats[param_name], esti_bits_dict[param_name] = simulate_fn(splats[param_name], step, choose_idx)
+            else:
+                new_splats[param_name] = splats[param_name] + 0.
+                esti_bits_dict[param_name] = None
+        
+        return new_splats, esti_bits_dict
 
-# to simulate what happens in gsplat 's PngCompression()
-def _min_max_quantization_16bit(param: Tensor) -> Tensor:
-    maxs = torch.amax(param, dim=0)
-    mins = torch.amin(param, dim=0)
+    def simulate_compression_means(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # out = torch.clamp(param, -5, 5)
+        # out = inverse_log_transform(log_transform(clamped_param))
+        
+        # return out, None
+        return torch.nn.Identity()(param), None
+        
+    def simulate_compression_quats(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # fake quantize
+        if step < 10_000:
+            fq_out_dict = fake_quantize_ste(param, self.bds["quats"][0], self.bds["quats"][1], 8, self.quantization_sim_type)
+        else:
+            fq_out_dict = fake_quantize_ste(param, self.bds["quats"][0], self.bds["quats"][1], self.q_bitwidth["quats"], self.quantization_sim_type)
+        
+        # entropy constraint
+        if step > self.entropy_steps["quats"] and self.entropy_model_enable and self.entropy_model_option["quats"]:
+            # import pdb; pdb.set_trace()
+            if choose_idx is not None:
+                esti_bits = self.entropy_models["quats"](fq_out_dict["output_value"][choose_idx], fq_out_dict["q_step"])
+            else:
+                esti_bits = self.entropy_models["quats"](fq_out_dict["output_value"], fq_out_dict["q_step"])
 
-    param_norm = (param - mins) / (maxs - mins)
-    q_step = 1 / (2**16 - 1)
-    q_param_norm = (((param_norm / q_step).round() * q_step) - param_norm).detach() + param_norm
+            return fq_out_dict["output_value"], esti_bits
+        else:
+            return fq_out_dict["output_value"], None
 
-    q_param = q_param_norm * (maxs - mins) + mins
+    
+    def simulate_compression_scales(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # fake quantize
+        if step < 10_000:
+            fq_out_dict = fake_quantize_ste(param, self.bds["scales"][0], self.bds["scales"][1], 8, self.quantization_sim_type)
+        else:
+            fq_out_dict = fake_quantize_ste(param, self.bds["scales"][0], self.bds["scales"][1], self.q_bitwidth["scales"], self.quantization_sim_type)
 
-    return q_param
+        # entropy constraint
+        if step > self.entropy_steps["scales"] and self.entropy_model_enable and self.entropy_model_option["scales"]:
+            # import pdb; pdb.set_trace()
+            # factorized model
+            if choose_idx is not None:
+                esti_bits = self.entropy_models["scales"](fq_out_dict["output_value"][choose_idx], fq_out_dict["q_step"])
+            else:
+                esti_bits = self.entropy_models["scales"](fq_out_dict["output_value"], fq_out_dict["q_step"])
 
-# to simulate what happens in gsplat 's PngCompression()
-def _min_max_quantization(param: Tensor) -> Tensor: # seems not working...
-    maxs = torch.amax(param, dim=0)
-    mins = torch.amin(param, dim=0)
+            # gaussian model
+            # mean = torch.mean(fq_out_dict["output_value"][choose_idx])
+            # std = torch.std(fq_out_dict["output_value"][choose_idx])
+            # esti_bits = self.entropy_models["scales"](fq_out_dict["output_value"][choose_idx], mean, std, fq_out_dict["q_step"])
 
-    param_norm = (param - mins) / (maxs - mins)
-    q_step = 1 / (2**8 - 1)
-    q_param_norm = (((param_norm / q_step).round() * q_step) - param_norm).detach() + param_norm
+            return fq_out_dict["output_value"], esti_bits
+        else:
+            return fq_out_dict["output_value"], None
 
-    q_param = q_param_norm * (maxs - mins) + mins
+    
+    def simulate_compression_opacities(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # fake quantize
+        fq_out_dict = fake_quantize_ste(param, self.bds["opacities"][0], self.bds["opacities"][1], 8, self.quantization_sim_type)
 
-    return q_param
+        # entropy constraint
+        if step > self.entropy_steps["opacities"] and self.entropy_model_enable and self.entropy_model_option["opacities"]:
+            fq_out_dict["output_value"] = fq_out_dict["output_value"].unsqueeze(1)
+            if choose_idx is not None:
+                esti_bits = self.entropy_models["opacities"](fq_out_dict["output_value"][choose_idx], fq_out_dict["q_step"])
+            else:
+                esti_bits = self.entropy_models["opacities"](fq_out_dict["output_value"], fq_out_dict["q_step"])
+            return fq_out_dict["output_value"].squeeze(1), esti_bits
+        else:
+            return fq_out_dict["output_value"], None
+
+
+    def simulate_compression_colors(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # fake quantize
+        if step < 10_000:
+            fq_out_dict = fake_quantize_ste(param, self.bds["colors"][0], self.bds["colors"][1], 8, self.quantization_sim_type)
+        else:
+            fq_out_dict = fake_quantize_ste(param, self.bds["colors"][0], self.bds["colors"][1], self.q_bitwidth["colors"], self.quantization_sim_type)
+        
+        # entropy constraint
+        if step > self.entropy_steps["colors"] and self.entropy_model_enable and self.entropy_model_option["colors"]:
+            fq_out_dict["output_value"] = fq_out_dict["output_value"]
+            if choose_idx is not None:
+                esti_bits = self.entropy_models["colors"](fq_out_dict["output_value"][choose_idx], fq_out_dict["q_step"])
+            else:
+                esti_bits = self.entropy_models["colors"](fq_out_dict["output_value"], fq_out_dict["q_step"])
+            return fq_out_dict["output_value"], esti_bits
+        else:
+            return fq_out_dict["output_value"], None
+    
 
-def _ste_quantization_for_quats(param: Tensor) -> Tensor:
-    return STE_quant_for_quats.apply(param, 8)
+    def simulate_compression_features_dir(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # fake quantize
+        if step < 10_000:
+            fq_out_dict = fake_quantize_ste(param, self.bds["features_dir"][0], self.bds["features_dir"][1], 8, self.quantization_sim_type)
+        else:
+            fq_out_dict = fake_quantize_ste(param, self.bds["features_dir"][0], self.bds["features_dir"][1], self.q_bitwidth["features_dir"], self.quantization_sim_type)
 
-def _ste_quantization_given_q_step(param: Tensor) -> Tensor:
-    return STE_multistep.apply(param, 0.001)
+        # entropy constraint
+        if step > self.entropy_steps["features_dir"] and self.entropy_model_enable and self.entropy_model_option["features_dir"]:
+            fq_out_dict["output_value"] = fq_out_dict["output_value"]
+            if choose_idx is not None:
+                esti_bits = self.entropy_models["features_dir"](fq_out_dict["output_value"][choose_idx], fq_out_dict["q_step"])
+            else:
+                esti_bits = self.entropy_models["features_dir"](fq_out_dict["output_value"], fq_out_dict["q_step"])
+            return fq_out_dict["output_value"], esti_bits
+        else:
+            return fq_out_dict["output_value"], None
+    
 
-def _ste_only(param: torch.nn.Parameter) -> torch.nn.Parameter:
-    return param
-    # return STE.apply(param)
-    # return (param.detach() -  param.detach()) + param # not working...
+    def simulate_compression_features_time(self, param: torch.nn.Parameter, step: int, choose_idx: torch.Tensor) -> Tensor:
+        # fake quantize
+        if step < 10_000:
+            fq_out_dict = fake_quantize_ste(param, self.bds["features_time"][0], self.bds["features_time"][1], 8, self.quantization_sim_type)
+        else:
+            fq_out_dict = fake_quantize_ste(param, self.bds["features_time"][0], self.bds["features_time"][1], self.q_bitwidth["features_time"], self.quantization_sim_type)
 
-def _add_noise_to_simulate_quantization(param: Tensor) -> Tensor:
-    return param + torch.empty_like(param).uniform_(-0.5, 0.5) * 0.001
+        # entropy constraint
+        if step > self.entropy_steps["features_time"] and self.entropy_model_enable and self.entropy_model_option["features_time"]:
+            fq_out_dict["output_value"] = fq_out_dict["output_value"]
+            if choose_idx is not None:
+                esti_bits = self.entropy_models["features_time"](fq_out_dict["output_value"][choose_idx], fq_out_dict["q_step"])
+            else:
+                esti_bits = self.entropy_models["features_time"](fq_out_dict["output_value"], fq_out_dict["q_step"])
+            return fq_out_dict["output_value"], esti_bits
+        else:
+            return fq_out_dict["output_value"], None