Merge branch 'dev_compression_simulation' into main

Author: JasonLSC

.gitignore

@@ -138,6 +138,7 @@ temp
 !examples/benchmarks/compression/results/
 .github/
 
+<<<<<<< HEAD
 setup_gscodec.py
 *.ipynb
 Readme_GSCodec.md 
@@ -146,6 +147,9 @@ temp
 stats
 figs
 
+=======
+gsc_utils
+>>>>>>> dev_compression_simulation
 examples/benchmarks/compression/mpeg_exp
 examples/benchmarks/compression/rebuttal
 examples/benchmarks/viewer

examples/benchmarks/compression/final_exp/mcmc_tt_sim.sh

@@ -1,7 +1,7 @@
 # ----------------- Training Setting-------------- #
 SCENE_DIR="data/tandt"
 # eval all 9 scenes for benchmarking
-SCENE_LIST="train truck" #  truck
+SCENE_LIST="truck" # train truck
 # SCENE_LIST="garden bicycle stump bonsai counter kitchen room treehill flowers"
 
 # # 0.36M GSs
@@ -41,14 +41,15 @@ run_single_scene() {
     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 \
-        --rd_lambda $RD_LAMBDA \
-        --shN_ada_mask_opt
+    # 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 \
+    #     --rd_lambda $RD_LAMBDA \
+    #     --shN_ada_mask_opt \
+    #     --compression entropy_coding
 
 
     # eval: use vgg for lpips to align with other benchmarks
@@ -57,8 +58,8 @@ run_single_scene() {
         --data_dir $SCENE_DIR/$SCENE/ \
         --result_dir $RESULT_DIR/$SCENE/ \
         --lpips_net vgg \
-        --compression png \
-        --ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt
+        --ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt \
+        --compression entropy_coding
 
 }
 # ----------------- Main Job --------------------- #
@@ -76,7 +77,7 @@ do
     SCENE_IDX=$((SCENE_IDX + 1))
     {
         run_single_scene ${GPU_LIST[$SCENE_IDX]} $SCENE
-    } &
+    } #&
 
 done
 

examples/benchmarks/compression/final_exp/mcmc_tt_sim_grad_thres.sh

@@ -0,0 +1,96 @@
+# ----------------- Training Setting-------------- #
+SCENE_DIR="data/tandt"
+# eval all 9 scenes for benchmarking
+SCENE_LIST="truck train" #  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_grad_thres"
+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 \
+        --rd_lambda $RD_LAMBDA \
+        --shN_ada_mask_opt --shN_ada_mask_strategy gradient
+
+
+    # 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 png \
+        --ckpt $RESULT_DIR/$SCENE/ckpts/ckpt_29999_rank0.pt
+    
+}
+# ----------------- Main Job --------------------- #
+
+
+
+# ----------------- Experiment Loop -------------- #
+GPU_LIST=(5 6)
+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