From 4ff628ea65acf9e3fce89514a7fed6c52b2fd6aa Mon Sep 17 00:00:00 2001
From: Friedrich Branschke <f.branschke@tu-berlin.de>
Date: Mon, 16 Dec 2024 15:30:43 +0100
Subject: [PATCH] fixed benchmark scripts for compression

---
 examples/benchmarks/compression/mcmc.sh    |  2 +-
 examples/benchmarks/compression/mcmc_db.sh | 48 ++++++++++++++++++++++
 examples/benchmarks/compression/mcmc_tt.sh |  4 +-
 3 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 examples/benchmarks/compression/mcmc_db.sh

diff --git a/examples/benchmarks/compression/mcmc.sh b/examples/benchmarks/compression/mcmc.sh
index bcc1fcfb0..29f221dae 100644
--- a/examples/benchmarks/compression/mcmc.sh
+++ b/examples/benchmarks/compression/mcmc.sh
@@ -43,7 +43,7 @@ done
 if command -v zip &> /dev/null
 then
     echo "Zipping results"
-    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR
+    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
 else
     echo "zip command not found, skipping zipping"
 fi
diff --git a/examples/benchmarks/compression/mcmc_db.sh b/examples/benchmarks/compression/mcmc_db.sh
new file mode 100644
index 000000000..716263fb0
--- /dev/null
+++ b/examples/benchmarks/compression/mcmc_db.sh
@@ -0,0 +1,48 @@
+SCENE_DIR="data/db"
+# eval all 9 scenes for benchmarking
+SCENE_LIST="playroom drjohnson"
+
+# # 0.36M GSs
+# RESULT_DIR="results/benchmark_db_mcmc_0_36M_png_compression"
+# CAP_MAX=360000
+
+# # 0.49M GSs
+# RESULT_DIR="results/benchmark_db_mcmc_0_49M_png_compression"
+# CAP_MAX=490000
+
+# 1M GSs
+RESULT_DIR="results/benchmark_db_mcmc_1M_png_compression"
+CAP_MAX=1000000
+
+# # 4M GSs
+# RESULT_DIR="results/benchmark_db_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 --scenes $SCENE_LIST
+else
+    echo "zip command not found, skipping zipping"
+fi
diff --git a/examples/benchmarks/compression/mcmc_tt.sh b/examples/benchmarks/compression/mcmc_tt.sh
index 054920929..34dba3ccd 100644
--- a/examples/benchmarks/compression/mcmc_tt.sh
+++ b/examples/benchmarks/compression/mcmc_tt.sh
@@ -42,7 +42,7 @@ done
 if command -v zip &> /dev/null
 then
     echo "Zipping results"
-    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST
+    python benchmarks/compression/summarize_stats.py --results_dir $RESULT_DIR --scenes $SCENE_LIST --scenes $SCENE_LIST
 else
     echo "zip command not found, skipping zipping"
-fi
\ No newline at end of file
+fi