From 265267e0e76ae0e1c9ae9eea39f3e475f66fe92c Mon Sep 17 00:00:00 2001
From: Ruilong Li <397653553@qq.com>
Date: Thu, 22 Aug 2024 03:54:21 +0000
Subject: [PATCH] tt compression script

---
 examples/benchmarks/compression/mcmc_tt.sh    | 48 +++++++++++++++++++
 .../compression/results/TanksAndTemples.csv   |  5 ++
 .../benchmarks/compression/summarize_stats.py | 15 ++----
 3 files changed, 56 insertions(+), 12 deletions(-)
 create mode 100644 examples/benchmarks/compression/mcmc_tt.sh
 create mode 100644 examples/benchmarks/compression/results/TanksAndTemples.csv

diff --git a/examples/benchmarks/compression/mcmc_tt.sh b/examples/benchmarks/compression/mcmc_tt.sh
new file mode 100644
index 000000000..054920929
--- /dev/null
+++ b/examples/benchmarks/compression/mcmc_tt.sh
@@ -0,0 +1,48 @@
+SCENE_DIR="data/tandt"
+# eval all 9 scenes for benchmarking
+SCENE_LIST="train 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/benchmark_tt_mcmc_1M_png_compression"
+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
\ No newline at end of file
diff --git a/examples/benchmarks/compression/results/TanksAndTemples.csv b/examples/benchmarks/compression/results/TanksAndTemples.csv
new file mode 100644
index 000000000..5845808d9
--- /dev/null
+++ b/examples/benchmarks/compression/results/TanksAndTemples.csv
@@ -0,0 +1,5 @@
+Submethod,PSNR,SSIM,LPIPS,Size [Bytes],#Gaussians
+,23.54,0.838,0.200,6875669,360000
+,23.62,0.845,0.188,8728572,490000
+-1.00M,24.03,0.857,0.163,16100628,1000000
+,24.47,0.872,0.132,58239022,4000000
diff --git a/examples/benchmarks/compression/summarize_stats.py b/examples/benchmarks/compression/summarize_stats.py
index d6aca0cb5..d11dbed6f 100644
--- a/examples/benchmarks/compression/summarize_stats.py
+++ b/examples/benchmarks/compression/summarize_stats.py
@@ -2,23 +2,14 @@
 import os
 import subprocess
 from collections import defaultdict
+from typing import List
 
 import numpy as np
 import tyro
 
 
-def main(results_dir: str = "results/benchmark_mcmc_0_36M_png_compression"):
-    scenes = [
-        "garden",
-        "bicycle",
-        "stump",
-        "bonsai",
-        "counter",
-        "kitchen",
-        "room",
-        "treehill",
-        "flowers",
-    ]
+def main(results_dir: str, scenes: List[str]):
+    print("scenes:", scenes)
     stage = "compress"
 
     summary = defaultdict(list)