feat: entropy coding from learned dist. (from Gaussian model)

Author: JasonLSC

.gitignore

@@ -133,5 +133,6 @@ Readme_GSCodec.md
 
 figs
 stats
+temp
 
 !examples/benchmarks/compression/results/

examples/benchmarks/compression/final_exp/mcmc_tt_sim_hash_grid.sh

@@ -0,0 +1,97 @@
+# ----------------- Training Setting-------------- #
+SCENE_DIR="data/tandt"
+# eval all 9 scenes for benchmarking
+SCENE_LIST="train truck" #  truck
+# SCENE_LIST="garden bicycle stump bonsai counter kitchen room treehill flowers"
+
+# # 0.36M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_0_36M_png_compression"
+# CAP_MAX=360000
+
+# # 0.49M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_tt_0_49M_png_compression"
+# CAP_MAX=490000
+
+# 1M GSs
+RESULT_DIR="results/Ours_TT_hash_grid"
+CAP_MAX=1000000
+
+# # 4M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_4M_png_compression"
+# CAP_MAX=4000000
+
+RD_LAMBDA=0.01
+
+# ----------------- Training Setting-------------- #
+
+# ----------------- Args ------------------------- #
+
+if [ ! -z "$1" ]; then
+    RD_LAMBDA="$1"
+    RESULT_DIR="results/Ours_TT_rd_lambda_${RD_LAMBDA}"
+fi
+
+# ----------------- Args ------------------------- #
+
+# ----------------- Main Job --------------------- #
+run_single_scene() {
+    local GPU_ID=$1
+    local SCENE=$2
+
+    echo "Running $SCENE on GPU: $GPU_ID"
+
+    # train without eval
+    CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --eval_steps -1 --disable_viewer --data_factor 1 \
+        --strategy.cap-max $CAP_MAX \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/ \
+        --compression_sim \
+        --entropy_model_opt --entropy_model_type gaussian_model \
+        --rd_lambda $RD_LAMBDA \
+        --shN_ada_mask_opt \
+        --compression entropy_coding 
+
+
+    # eval: use vgg for lpips to align with other benchmarks
+    CUDA_VISIBLE_DEVICES=$GPU_ID python simple_trainer.py mcmc --disable_viewer --data_factor 1 \
+        --strategy.cap-max $CAP_MAX \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/ \
+        --lpips_net vgg \
+        --compression entropy_coding --entropy_model_type gaussian_model \
+        --ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt
+    
+}
+# ----------------- Main Job --------------------- #
+
+
+
+# ----------------- Experiment Loop -------------- #
+# GPU_LIST=(5 7)
+# GPU_COUNT=${#GPU_LIST[@]}
+
+# SCENE_IDX=-1
+
+# for SCENE in $SCENE_LIST;
+# do
+#     SCENE_IDX=$((SCENE_IDX + 1))
+#     {
+#         run_single_scene ${GPU_LIST[$SCENE_IDX]} $SCENE
+#     } &
+
+# done
+
+# ----------------- Experiment Loop -------------- #
+
+# Wait for finishing the jobs across all scenes 
+wait
+echo "All scenes finished."
+
+# Zip the compressed files and summarize the stats
+if command -v zip &> /dev/null
+then
+    echo "Zipping results"
+    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
+else
+    echo "zip command not found, skipping zipping"
+fi

examples/benchmarks/compression/mcmc_tt_ans_debug.sh

@@ -0,0 +1,48 @@
+SCENE_DIR="data/tandt"
+# eval all 2 scenes for benchmarking
+SCENE_LIST="train truck" # truck
+
+# # 0.36M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_0_36M_png_compression"
+# CAP_MAX=360000
+
+# # 0.49M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_tt_0_49M_png_compression"
+# CAP_MAX=490000
+
+# 1M GSs
+RESULT_DIR="results/Ours_TT_rd_lambda_0.002_qualitative"
+CAP_MAX=1000000
+
+# # 4M GSs
+# RESULT_DIR="results/benchmark_tt_mcmc_4M_png_compression"
+# CAP_MAX=4000000
+
+for SCENE in $SCENE_LIST;
+do
+    echo "Running $SCENE"
+
+    # train without eval
+    # CUDA_VISIBLE_DEVICES=0 python simple_trainer.py mcmc --eval_steps -1 --disable_viewer --data_factor 1 \
+    #     --strategy.cap-max $CAP_MAX \
+    #     --data_dir $SCENE_DIR/$SCENE/ \
+    #     --result_dir $RESULT_DIR/$SCENE/
+
+    # eval: use vgg for lpips to align with other benchmarks
+    CUDA_VISIBLE_DEVICES=0 python simple_trainer.py mcmc --disable_viewer --data_factor 1 \
+        --strategy.cap-max $CAP_MAX \
+        --data_dir $SCENE_DIR/$SCENE/ \
+        --result_dir $RESULT_DIR/$SCENE/ \
+        --lpips_net vgg \
+        --compression png \
+        --ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt
+done
+
+# Zip the compressed files and summarize the stats
+if command -v zip &> /dev/null
+then
+    echo "Zipping results"
+    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
+else
+    echo "zip command not found, skipping zipping"
+fi

examples/simple_trainer.py

@@ -30,6 +30,7 @@
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 from typing_extensions import Literal, assert_never
 from gsplat import strategy
+from gsplat.compression.entropy_coding_compression import EntropyCodingCompression
 from gsplat.compression_simulation import simulation
 from utils import AppearanceOptModule, CameraOptModule, knn, rgb_to_sh, set_random_seed
 from lib_bilagrid import (
@@ -44,6 +45,7 @@
 from gsplat.rendering import rasterization
 from gsplat.strategy import DefaultStrategy, MCMCStrategy
 from gsplat.compression_simulation import CompressionSimulation
+from gsplat.compression_simulation.entropy_model import Entropy_factorized_optimized_refactor, Entropy_gaussian
 
 class ProfilerConfig:
     def __init__(self):
@@ -92,7 +94,7 @@ class Config:
     # Path to the .pt files. If provide, it will skip training and run evaluation only.
     ckpt: Optional[List[str]] = None
     # Name of compression strategy to use
-    compression: Optional[Literal["png"]] = None
+    compression: Optional[Literal["png", "entropy_coding"]] = None
 
     # Enable profiler
     profiler_enabled: bool = False
@@ -118,10 +120,10 @@ class Config:
                                                                    "shN": 10_000})
     # gaussian model:
     # entropy_steps: Dict[str, int] = field(default_factory=lambda: {"means": -1, 
-    #                                                                "quats": -1, 
+    #                                                                "quats": 10_000, 
     #                                                                "scales": 10_000, 
-    #                                                                "opacities": -1, 
-    #                                                                "sh0": -1, 
+    #                                                                "opacities": 10_000, 
+    #                                                                "sh0": 20_000, 
     #                                                                "shN": -1})
 
     # Enable shN adaptive mask
@@ -420,6 +422,8 @@ def __init__(
         if cfg.compression is not None:
             if cfg.compression == "png":
                 self.compression_method = PngCompression()
+            elif  cfg.compression == "entropy_coding":
+                self.compression_method = EntropyCodingCompression()
             else:
                 raise ValueError(f"Unknown compression strategy: {cfg.compression}")
 
@@ -897,6 +901,11 @@ def train(self):
 
                     if cfg.shN_ada_mask_opt and step > cfg.ada_mask_steps:
                         data["shN_ada_mask"] = shN_ada_mask
+                    
+                    if cfg.compression_sim and cfg.entropy_model_opt and cfg.compression == "entropy_coding":
+                        for name, entropy_model in self.compression_sim_method.entropy_models.items():
+                            if entropy_model is not None:
+                                data[name+"_entropy_model"] = entropy_model.state_dict()
 
                     torch.save(
                         data, f"{self.ckpt_dir}/ckpt_{step}_rank{self.world_rank}.pt"
@@ -1032,8 +1041,9 @@ def eval(self, step: int, stage: str = "val"):
             canvas_list = [pixels, colors]
 
             if world_rank == 0:
-                # write images
-                canvas = torch.cat(canvas_list, dim=2).squeeze(0).cpu().numpy()
+                # write images 
+                # canvas = torch.cat(canvas_list, dim=2).squeeze(0).cpu().numpy() # side by side
+                canvas = canvas_list[1].squeeze(0).cpu().numpy() # signle image
                 canvas = (canvas * 255).astype(np.uint8)
                 imageio.imwrite(
                     f"{self.render_dir}/{stage}_step{step}_{i:04d}.png",
@@ -1074,16 +1084,18 @@ def eval(self, step: int, stage: str = "val"):
             self.writer.flush()
 
     @torch.no_grad()
-    def render_traj(self, step: int):
+    def render_traj(self, step: int, stage: str = "val"):
         """Entry for trajectory rendering."""
         print("Running trajectory rendering...")
         cfg = self.cfg
         device = self.device
 
-        camtoworlds_all = self.parser.camtoworlds[5:-5]
+        num_imgs = len(self.parser.camtoworlds)
+
+        camtoworlds_all = self.parser.camtoworlds[: num_imgs//2]
         if cfg.render_traj_path == "interp":
             camtoworlds_all = generate_interpolated_path(
-                camtoworlds_all, 1
+                camtoworlds_all, 6 #1
             )  # [N, 3, 4]
         elif cfg.render_traj_path == "ellipse":
             height = camtoworlds_all[:, 2, 3].mean()
@@ -1118,7 +1130,7 @@ def render_traj(self, step: int):
         # save to video
         video_dir = f"{cfg.result_dir}/videos"
         os.makedirs(video_dir, exist_ok=True)
-        writer = imageio.get_writer(f"{video_dir}/traj_{step}.mp4", fps=30)
+        writer = imageio.get_writer(f"{video_dir}/{stage}_traj_{step}.mp4", fps=30)
         for i in tqdm.trange(len(camtoworlds_all), desc="Rendering trajectory"):
             camtoworlds = camtoworlds_all[i : i + 1]
             Ks = K[None]
@@ -1139,11 +1151,12 @@ def render_traj(self, step: int):
             canvas_list = [colors, depths.repeat(1, 1, 1, 3)]
 
             # write images
-            canvas = torch.cat(canvas_list, dim=2).squeeze(0).cpu().numpy()
+            # canvas = torch.cat(canvas_list, dim=2).squeeze(0).cpu().numpy()
+            canvas = canvas_list[0].squeeze(0).cpu().numpy()
             canvas = (canvas * 255).astype(np.uint8)
             writer.append_data(canvas)
         writer.close()
-        print(f"Video saved to {video_dir}/traj_{step}.mp4")
+        print(f"Video saved to {video_dir}/{stage}_traj_{step}.mp4")
 
     @torch.no_grad()
     def run_compression(self, step: int):
@@ -1156,8 +1169,12 @@ def run_compression(self, step: int):
         # import pdb; pdb.set_trace()
         self.run_param_distribution_vis(self.splats, save_dir=f"{cfg.result_dir}/visualization/raw")
 
-        self.compression_method.compress(compress_dir, self.splats)
-        # self.run_param_distribution_vis(self.splats, save_dir=f"{cfg.result_dir}/visualization/log_transform")
+        if isinstance(self.compression_method, PngCompression):
+            self.compression_method.compress(compress_dir, self.splats)
+        elif isinstance(self.compression_method, EntropyCodingCompression):
+            self.compression_method.compress(compress_dir, self.splats, self.entropy_models)
+        else:
+            raise NotImplementedError(f"The compression method is not implemented yet.")
 
         # evaluate compression
         splats_c = self.compression_method.decompress(compress_dir)
@@ -1167,6 +1184,7 @@ def run_compression(self, step: int):
         for k in splats_c.keys():
             self.splats[k].data = splats_c[k].to(self.device)
         self.eval(step=step, stage="compress")
+        self.render_traj(step=step, stage="compress")
 
     @torch.no_grad()
     def run_param_distribution_vis(self, param_dict: Dict[str, Tensor], save_dir: str):
@@ -1199,6 +1217,21 @@ def run_param_distribution_vis(self, param_dict: Dict[str, Tensor], save_dir: st
             plt.close()
 
         print(f"Histograms saved in '{save_dir}' directory.")
+    
+    def load_entropy_model_from_ckpt(self, ckpt: Dict, entropy_model_type: str):
+        self.entropy_models = {}
+        for name, value in ckpt.items():
+            if "_entropy_model" in name:
+                attr_name = name[:(len(name) - len("_entropy_model"))]
+                num_ch = ckpt["splats"][attr_name].shape[-1]
+                if entropy_model_type == "factorized_model":
+                    # TODO
+                    pass
+                elif entropy_model_type == "gaussian_model":
+                    entropy_model = Entropy_gaussian(channel=num_ch)
+                
+                entropy_model.load_state_dict(value)
+                self.entropy_models[attr_name] = entropy_model
 
     @torch.no_grad()
     def _viewer_render_fn(
@@ -1242,6 +1275,8 @@ def main(local_rank: int, world_rank, world_size: int, cfg: Config):
         runner.eval(step=step)
         runner.render_traj(step=step)
         if cfg.compression is not None:
+            if cfg.compression == "entropy_coding":
+                runner.load_entropy_model_from_ckpt(ckpts[0], cfg.entropy_model_type)
             runner.run_compression(step=step)
     else:
         runner.train()

gsplat/compression/__init__.py

@@ -1 +1,2 @@
 from .png_compression import PngCompression
+from .entropy_coding_compression import EntropyCodingCompression