diff --git a/INSTALL.sh b/INSTALL.sh new file mode 100644 index 0000000000000000000000000000000000000000..d8f6c2233cea51f07c06fb5953f689b1abac0dd6 --- /dev/null +++ b/INSTALL.sh @@ -0,0 +1,105 @@ +#!/bin/bash +# Deployment script for splatatlas_full_patch.tar.gz +# Run from /root/autodl-tmp/SplatAtlas/ +# +# What this does: +# 1. Sanity-check we're in the right directory +# 2. Back up the 4 files that will be overwritten +# 3. Extract the patch (methods/, configs/, run_*.py, INSTALL.sh) +# 4. Auto-generate pixelgssgf_benchmark.yaml and sgfgs_benchmark.yaml +# from absgs template if they don't exist (with method_name swapped) +# 5. Run verification grep checks + +set -e + +EXPECTED_ROOT="/root/autodl-tmp/SplatAtlas" +if [ "$(pwd)" != "$EXPECTED_ROOT" ]; then + echo "ERROR: please cd to $EXPECTED_ROOT first" + echo " (you are in $(pwd))" + exit 1 +fi + +PATCH=splatatlas_full_patch.tar.gz +if [ ! -f "$PATCH" ]; then + echo "ERROR: $PATCH not found in $(pwd)" + exit 1 +fi + +# ---- 1. Backup the 4 modified files ---- +BACKUP_DIR=".backup_$(date +%Y%m%d_%H%M%S)" +mkdir -p "$BACKUP_DIR" +for f in methods/wrapper_absgs.py methods/wrapper_pixelgs.py \ + methods/wrapper_pixelgssgf.py methods/wrapper_sgfgs.py; do + if [ -f "$f" ]; then + cp "$f" "$BACKUP_DIR/" + fi +done +echo "[1/4] backed up 4 modified wrappers → $BACKUP_DIR/" + +# ---- 2. Extract patch ---- +tar -xzf "$PATCH" +echo "[2/4] patch extracted" + +# ---- 3. Auto-generate the two configs that might be missing ---- +for missing in pixelgssgf sgfgs; do + target="configs/${missing}_benchmark.yaml" + if [ ! -f "$target" ]; then + if [ -f configs/absgs_benchmark.yaml ]; then + sed "s/^method_name: absgs/method_name: ${missing}/" \ + configs/absgs_benchmark.yaml > "$target" + echo "[3/4] generated missing $target from absgs template" + else + echo "[3/4] WARNING: $target missing AND no absgs_benchmark.yaml to clone from" + fi + fi +done + +# ---- 4. Verification ---- +echo "[4/4] verification:" +echo "" + +echo "--- SGF residue in wrapper_pixelgs.py (should be empty) ---" +grep -n "SGFController\|guided_l1_loss\|guided_ssim_loss" methods/wrapper_pixelgs.py \ + && echo " FAIL: residue found" || echo " PASS" +echo "" + +echo "--- pixels kept in wrapper_pixelgs.py (should have 2 lines) ---" +grep -n "pixels" methods/wrapper_pixelgs.py | grep -E "render_pkg|add_densification" || echo " FAIL" +echo "" + +echo "--- new wrappers registered ---" +grep -n "register_method" methods/wrapper_edgeloss.py methods/wrapper_vanillasgf.py +echo "" + +echo "--- timing in all 6 wrappers ---" +for f in absgs pixelgs pixelgssgf sgfgs vanillasgf edgeloss; do + if [ -f "methods/wrapper_${f}.py" ]; then + n_time=$(grep -c "iter_time_ms\|render_ms\|vram_allocated" methods/wrapper_${f}.py) + echo " wrapper_${f}.py: $n_time timing references" + fi +done +echo "" + +echo "--- configs present ---" +ls -1 configs/*_benchmark.yaml +echo "" + +echo "--- run scripts present ---" +ls -1 run_*.py +echo "" + +echo "--- master execution plan dependency check ---" +for need in run_vanillasgf.py run_absgs.py run_pixelgs.py run_edgeloss.py run_vanilla.py; do + if [ -f "$need" ]; then + echo " OK $need" + else + echo " MISS $need (run_all.py will skip this method)" + fi +done +echo "" + +echo "DONE. To run a single method: python run_.py" +echo " To run all methods: python run_all.py" +echo " To dry-run plan: python run_all.py --dry-run" +echo "" +echo "Backup of pre-patch wrappers: $BACKUP_DIR/" diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/configs/2dgs_benchmark.yaml b/configs/2dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e036cddc060d41586977c20b736c35174875e7f2 --- /dev/null +++ b/configs/2dgs_benchmark.yaml @@ -0,0 +1,17 @@ +method_name: 2dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - church + - courthouse + - ignatius + - meetingroom + - truck diff --git a/configs/2dgssgf_tnt_benchmark.yaml b/configs/2dgssgf_tnt_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ccb185bdaf0c2374653e99bdcb9adb1d18e9acca --- /dev/null +++ b/configs/2dgssgf_tnt_benchmark.yaml @@ -0,0 +1,17 @@ +method_name: 2dgssgf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - church + - courthouse + - ignatius + - meetingroom + - truck diff --git a/configs/3dgs_dr_benchmark.yaml b/configs/3dgs_dr_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..edd599b9c3dfe280a37be09231f428e1a36ba2ab --- /dev/null +++ b/configs/3dgs_dr_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: 3dgs_dr +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/3dgsmcmc_benchmark.yaml b/configs/3dgsmcmc_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49182f908a849993fe8c6c07b837a2bbda7e3a59 --- /dev/null +++ b/configs/3dgsmcmc_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: 3dgsmcmc +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/3dgut_benchmark.yaml b/configs/3dgut_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..616d8d4b0ed9f1d32e6a7739b954ad5d39c76298 --- /dev/null +++ b/configs/3dgut_benchmark.yaml @@ -0,0 +1,12 @@ +method_name: 3dgut +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Mic + - Ship diff --git a/configs/aaags_benchmark.yaml b/configs/aaags_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c7264e14403f3df821031bc7587d6136285ecb0 --- /dev/null +++ b/configs/aaags_benchmark.yaml @@ -0,0 +1,16 @@ +method_name: aaags +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - caterpillar + - truck + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/absgs_benchmark.yaml b/configs/absgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b08a153b07cce2cc81e405f91e38b31a47c27d26 --- /dev/null +++ b/configs/absgs_benchmark.yaml @@ -0,0 +1,41 @@ +method_name: absgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - auditorium + - ballroom + - courtroom + - family + - francis + - horse + - lighthouse + - m60 + - museum + - palace + - panther + - playground + - temple + - train +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill diff --git a/configs/absgssgf_benchmark.yaml b/configs/absgssgf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7516f4aeea2d164f8178dc99207395dd806f9a0e --- /dev/null +++ b/configs/absgssgf_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: absgssgf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/analyticsplatting_benchmark.yaml b/configs/analyticsplatting_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3d253212894e2662acf055f74435eaac60c5c744 --- /dev/null +++ b/configs/analyticsplatting_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: analyticsplatting +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/analyticsplatting_benchmark.yaml.four_scenes b/configs/analyticsplatting_benchmark.yaml.four_scenes new file mode 100644 index 0000000000000000000000000000000000000000..d8114f6ebf594f1f1272bf90051867c382d70eab --- /dev/null +++ b/configs/analyticsplatting_benchmark.yaml.four_scenes @@ -0,0 +1,11 @@ +method_name: analyticsplatting +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - ignatius diff --git a/configs/atomgs_benchmark.yaml b/configs/atomgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..91315c4aea85a2943590a48afa18b4341de10a8b --- /dev/null +++ b/configs/atomgs_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: atomgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - ignatius diff --git a/configs/bags_benchmark.yaml b/configs/bags_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3ff6f08cfce43341cc92e4cdf1261be30f92e53d --- /dev/null +++ b/configs/bags_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: bags +global_settings: + iterations: 60000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/c3dgs_benchmark.yaml b/configs/c3dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7ad390c7390675a4c97f1ed1637bd0dd10a3258c --- /dev/null +++ b/configs/c3dgs_benchmark.yaml @@ -0,0 +1,10 @@ +method_name: "c3dgs" +global_settings: + iterations: 5000 + seeds: [0] + +datasets: + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Lego"] diff --git a/configs/cdcgs_benchmark.yaml b/configs/cdcgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..27f9eb19d1eb01229d532b851549b1d8ae9a4560 --- /dev/null +++ b/configs/cdcgs_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: cdcgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/coadaptgs_benchmark.yaml b/configs/coadaptgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d3074a29a05de22813028a0134b690abe36e78ba --- /dev/null +++ b/configs/coadaptgs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: coadaptgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/compact3dgs_benchmark.yaml b/configs/compact3dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3c5a86f2fb32db4a25d58e09d7624bfc4d889650 --- /dev/null +++ b/configs/compact3dgs_benchmark.yaml @@ -0,0 +1,25 @@ +method_name: "compact3dgs" +global_settings: + iterations: 30000 + seeds: [0] +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/compgs_benchmark.yaml b/configs/compgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4e6ace74cda9e2e2924a20a1841ad451b9ceec05 --- /dev/null +++ b/configs/compgs_benchmark.yaml @@ -0,0 +1,25 @@ +method_name: "compgs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["truck", "train"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "garden"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Lego"] diff --git a/configs/conegs_benchmark.yaml b/configs/conegs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..40528515308b9dd6b0e05dc9401adc57e9a4d148 --- /dev/null +++ b/configs/conegs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: conegs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/contextgs_benchmark.yaml b/configs/contextgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a7a5c64c2f502ec2af07850b524be62bbd4082da --- /dev/null +++ b/configs/contextgs_benchmark.yaml @@ -0,0 +1,25 @@ +method_name: "contextgs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] diff --git a/configs/drkgs_benchmark.yaml b/configs/drkgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fa12ad512e6ec526389305cc8fb570ed6e8b952c --- /dev/null +++ b/configs/drkgs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: drkgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/eagles_benchmark.yaml b/configs/eagles_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8bb6933cf898885933354b200dbf3d59306fd40c --- /dev/null +++ b/configs/eagles_benchmark.yaml @@ -0,0 +1,20 @@ +method_name: "eagles" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] diff --git a/configs/edgeloss_benchmark.yaml b/configs/edgeloss_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ff6a9792d7b39a5f4ae32cc661e097133cc0a8d5 --- /dev/null +++ b/configs/edgeloss_benchmark.yaml @@ -0,0 +1,41 @@ +method_name: edgeloss +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - auditorium + - ballroom + - courtroom + - family + - francis + - horse + - lighthouse + - m60 + - museum + - palace + - panther + - playground + - temple + - train +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill diff --git a/configs/erankgs_benchmark.yaml b/configs/erankgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b0de07664912d926749423e0d7457f6c4ae564d3 --- /dev/null +++ b/configs/erankgs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: erankgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/fds_benchmark.yaml b/configs/fds_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9bbc101b64bd0f3cfd759497ac2c70533e2d4eb2 --- /dev/null +++ b/configs/fds_benchmark.yaml @@ -0,0 +1,10 @@ +method_name: "fds" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter"] diff --git a/configs/fdsgs_benchmark.yaml b/configs/fdsgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e8489e57729143363dbe9b636ad5ef81e63804b4 --- /dev/null +++ b/configs/fdsgs_benchmark.yaml @@ -0,0 +1,20 @@ +method_name: "fdsgs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["truck", "train"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["room", "bonsai"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["garden", "bicycle"] diff --git a/configs/flod_benchmark.yaml b/configs/flod_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e028addffcf93ea842f9d41d932acc2fd49b7091 --- /dev/null +++ b/configs/flod_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "flod" +global_settings: + iterations: 100000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/gaussianfocus_benchmark.yaml b/configs/gaussianfocus_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6727455df995744a9e50be3bd10d974f0340d006 --- /dev/null +++ b/configs/gaussianfocus_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "gaussianfocus" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/gaussianpro_benchmark.yaml b/configs/gaussianpro_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9002881d28e5c5622989e5fe6e9f26b38369717a --- /dev/null +++ b/configs/gaussianpro_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: gaussianpro +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/gaussianshader_benchmark.yaml b/configs/gaussianshader_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2aa3f5a9a9a8bfdc6f78230e8e090d2a28638a3c --- /dev/null +++ b/configs/gaussianshader_benchmark.yaml @@ -0,0 +1,35 @@ +method_name: "gaussianshader" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/geogaussian_benchmark.yaml b/configs/geogaussian_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..604567f10a01d59b4fbc33b1fbcbf14a982d656c --- /dev/null +++ b/configs/geogaussian_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "geogaussian" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/ges_benchmark.yaml b/configs/ges_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f8be5a0c9504876bdcfcb131bdcd94ee74fc326a --- /dev/null +++ b/configs/ges_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: ges +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/ghap_benchmark.yaml b/configs/ghap_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6f1f4243c856e0427645a40b370610f6d108535a --- /dev/null +++ b/configs/ghap_benchmark.yaml @@ -0,0 +1,17 @@ +method_name: ghap +global_settings: + iterations: 30000 + seeds: + - 0 + hyperparams: + sampling_iter: 15000 + sampling_ratio: 0.1 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/gigs_benchmark.yaml b/configs/gigs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..69865d9455d81ead84589793eb1cbdaab5dc279c --- /dev/null +++ b/configs/gigs_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "gigs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/gof_benchmark.yaml b/configs/gof_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..64bb4abd6ceb8858b4148c4c10dca8d9137f64dc --- /dev/null +++ b/configs/gof_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: gof +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/gsdf_benchmark.yaml b/configs/gsdf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dab144b71294bd5ba8d50bcef235d6f15db55d34 --- /dev/null +++ b/configs/gsdf_benchmark.yaml @@ -0,0 +1,12 @@ +method_name: gsdf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - caterpillar + - truck diff --git a/configs/gsgn_benchmark.yaml b/configs/gsgn_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8025ccdd92cf041cc1e6d4736ca0fc3f6d487e53 --- /dev/null +++ b/configs/gsgn_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "gsgn" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/gslpm_benchmark.yaml b/configs/gslpm_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..27d529ca95eae50d3d0ffabeb20b13531b5a9c85 --- /dev/null +++ b/configs/gslpm_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: gslpm +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/gspull_benchmark.yaml b/configs/gspull_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b08ef39800b22dfe347b820b8c1b3d1c248b33bf --- /dev/null +++ b/configs/gspull_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: gspull +global_settings: + iterations: 15000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/gsshader_benchmark.yaml b/configs/gsshader_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3422b7b4440faf8232c3610902b099acdeb5eda3 --- /dev/null +++ b/configs/gsshader_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "gsshader" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/hac_plus_benchmark.yaml b/configs/hac_plus_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..63ac352de3daa031db7add9b427daf031f75c29b --- /dev/null +++ b/configs/hac_plus_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "hac_plus" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/hogs_benchmark.yaml b/configs/hogs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..49333fcf1aa8cb9840ac7bd983577e8ebbc64b1f --- /dev/null +++ b/configs/hogs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: hogs +global_settings: + iterations: 50000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/improvingadc_benchmark.yaml b/configs/improvingadc_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fe70d32b417af6ce8c191a6e1fb1b24bef94e379 --- /dev/null +++ b/configs/improvingadc_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "improvingadc" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/lapisgs_benchmark.yaml b/configs/lapisgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..30ed3dca5361e46796a3192351ca7d2f0fb7c844 --- /dev/null +++ b/configs/lapisgs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: lapisgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/lightgaussian_benchmark.yaml b/configs/lightgaussian_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f85ea5e88a063c4561878fab08777af5616a7897 --- /dev/null +++ b/configs/lightgaussian_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: lightgaussian +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/lod_gs_benchmark.yaml b/configs/lod_gs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9f617ad3e7500bcc0c70550c355a41ad1b8f71f1 --- /dev/null +++ b/configs/lod_gs_benchmark.yaml @@ -0,0 +1,15 @@ +method_name: lod_gs +global_settings: + iterations: 30000 + cap_gaussians: 1500000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/lp3dgs_benchmark.yaml b/configs/lp3dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8621e852d63f7780895a9e8c2264318fe55aad9c --- /dev/null +++ b/configs/lp3dgs_benchmark.yaml @@ -0,0 +1,20 @@ +method_name: "lp3dgs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["truck"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["garden"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["DrJohnson"] diff --git a/configs/mass_production.yaml b/configs/mass_production.yaml new file mode 100644 index 0000000000000000000000000000000000000000..04c89ac01d1c0d265fdfe9321707a3af3ba416e2 --- /dev/null +++ b/configs/mass_production.yaml @@ -0,0 +1,42 @@ +global: + iterations: 30000 + seeds: [0] # 现阶段先跑 Seed 0 快速拿到全局结论 + +tasks: + # 1. Deep Blending 室内场景先测 (只取最经典的 2 个) + - method: "tamgs" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + scenes: + - "Playroom" + - "DrJohnson" + + # 2. Mip-NeRF 360 无界室外场景 (9 个) + - method: "tamgs" + base_path: "/root/autodl-tmp/dataset/360" + scenes: + - "bicycle" + - "bonsai" + - "counter" + - "flowers" + - "garden" + - "kitchen" + - "room" + - "stump" + - "treehill" + + # 3. Tanks and Temples 大尺度场景 (12 个 NerfBaselines 钦点场景) + - method: "tamgs" + base_path: "/root/autodl-tmp/dataset/tnt" + scenes: + - "barn" + - "caterpillar" + - "truck" + - "playground" + - "lighthouse" + - "train" + - "auditorium" + - "ballroom" + - "courtroom" + - "museum" + - "palace" + - "temple" diff --git a/configs/matrix_cat1_cat5.yaml b/configs/matrix_cat1_cat5.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b1c0a2481e79fb13373075aa1ac337ba9ba11a4b --- /dev/null +++ b/configs/matrix_cat1_cat5.yaml @@ -0,0 +1,16 @@ +global_settings: + iterations: 30000 + save_iterations: [7000, 30000] + +tasks: + - method: "vanilla_3dgs" + scene: "cat1_shiny_toaster" + source_path: "/root/autodl-tmp/dataset/degradations/shiny/toaster" + + - method: "tamgs" + scene: "cat1_shiny_toaster" + source_path: "/root/autodl-tmp/dataset/degradations/shiny/toaster" + + - method: "vanilla_3dgs" + scene: "cat4_sparse_3views" + source_path: "/root/autodl-tmp/dataset/degradations/sparse/bonsai_3v" diff --git a/configs/minisplatting_benchmark.yaml b/configs/minisplatting_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a7e2a3461e7bd13487e770a7165354c798cd7cd2 --- /dev/null +++ b/configs/minisplatting_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: minisplatting +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/minisplattingsgf_benchmark.yaml b/configs/minisplattingsgf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2d84d467467306d94d0caffd3d16f15db3a7af92 --- /dev/null +++ b/configs/minisplattingsgf_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: minisplattingsgf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/mipsplatting_benchmark.yaml b/configs/mipsplatting_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..91527a2f3346019bf48d40349540485729956a1d --- /dev/null +++ b/configs/mipsplatting_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: mipsplatting +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/mipsplattingsgf_benchmark.yaml b/configs/mipsplattingsgf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7ece541bd348a50d3d7af4e0dce1f100c3bb9112 --- /dev/null +++ b/configs/mipsplattingsgf_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "mipsplattingsgf" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple", "family", "francis", "horse", "m60", "panther"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/ngs_benchmark.yaml b/configs/ngs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7e05957cb81cb393ff4676b651acbd42d79a2c08 --- /dev/null +++ b/configs/ngs_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: ngs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/octree_gs_benchmark.yaml b/configs/octree_gs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..03ce0b9c4350d8672d27e13ab4462bf15d1ef55a --- /dev/null +++ b/configs/octree_gs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: octree_gs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/op43dgs_benchmark.yaml b/configs/op43dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..da98e4277c7680b7a33f5bc19c4a0147d04dd5db --- /dev/null +++ b/configs/op43dgs_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "op43dgs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/opti3dgs_benchmark.yaml b/configs/opti3dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3a6b41ed2ce4357d0edd26304abc7bb5ade451ac --- /dev/null +++ b/configs/opti3dgs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: opti3dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/patch_aaags_benchmark.yaml b/configs/patch_aaags_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0c7264e14403f3df821031bc7587d6136285ecb0 --- /dev/null +++ b/configs/patch_aaags_benchmark.yaml @@ -0,0 +1,16 @@ +method_name: aaags +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - caterpillar + - truck + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/patch_analyticsplatting_benchmark.yaml b/configs/patch_analyticsplatting_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f169e90dfcc1387afc5d656b8183119409a7f7bd --- /dev/null +++ b/configs/patch_analyticsplatting_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: analyticsplatting +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/patch_atomgs_benchmark.yaml b/configs/patch_atomgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..91315c4aea85a2943590a48afa18b4341de10a8b --- /dev/null +++ b/configs/patch_atomgs_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: atomgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - ignatius diff --git a/configs/patch_gsdf_benchmark.yaml b/configs/patch_gsdf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dab144b71294bd5ba8d50bcef235d6f15db55d34 --- /dev/null +++ b/configs/patch_gsdf_benchmark.yaml @@ -0,0 +1,12 @@ +method_name: gsdf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - caterpillar + - truck diff --git a/configs/patch_rdogaussian_benchmark.yaml b/configs/patch_rdogaussian_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fde42ed88de2833e2fe888e5f503e7ce0c458004 --- /dev/null +++ b/configs/patch_rdogaussian_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: rdogaussian +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church diff --git a/configs/patch_scaffoldgs_benchmark.yaml b/configs/patch_scaffoldgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8c75a96640959adb2a9123c6a2726a3c86e045d5 --- /dev/null +++ b/configs/patch_scaffoldgs_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: scaffoldgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - courthouse diff --git a/configs/patch_steepgs_benchmark.yaml b/configs/patch_steepgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bc93debf34aecab08798790f9e859233bd1a43d2 --- /dev/null +++ b/configs/patch_steepgs_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: steepgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - ignatius diff --git a/configs/patch_sugar_benchmark.yaml b/configs/patch_sugar_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0aada8485b0f2eee2dd179eefbb347d8913498e3 --- /dev/null +++ b/configs/patch_sugar_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: sugar +global_settings: + iterations: 15000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church diff --git a/configs/patch_tamgs_benchmark.yaml b/configs/patch_tamgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dd63d41c01d372f413d931cbebe2e66067f22b67 --- /dev/null +++ b/configs/patch_tamgs_benchmark.yaml @@ -0,0 +1,12 @@ +method_name: tamgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - caterpillar + - truck diff --git a/configs/patch_vanilla_benchmark.yaml b/configs/patch_vanilla_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f538970e7ac456b0d2ab80156be3db3fea511b93 --- /dev/null +++ b/configs/patch_vanilla_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: vanilla_3dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - courthouse diff --git a/configs/pgsr_benchmark.yaml b/configs/pgsr_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..723054f4b760df26f6917e30675b118a91d631ba --- /dev/null +++ b/configs/pgsr_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: pgsr +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - meetingroom diff --git a/configs/pgsrsgf_tnt_benchmark.yaml b/configs/pgsrsgf_tnt_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..da19391424ea9665581e828a06171e9cee81a87b --- /dev/null +++ b/configs/pgsrsgf_tnt_benchmark.yaml @@ -0,0 +1,17 @@ +method_name: pgsrsgf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - church + - courthouse + - ignatius + - meetingroom + - truck diff --git a/configs/pixelgs_benchmark.yaml b/configs/pixelgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..71d0b8806fc88ba6ab8ad2d9d886e04edf091709 --- /dev/null +++ b/configs/pixelgs_benchmark.yaml @@ -0,0 +1,41 @@ +method_name: pixelgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - auditorium + - ballroom + - courtroom + - family + - francis + - horse + - lighthouse + - m60 + - museum + - palace + - panther + - playground + - temple + - train +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill diff --git a/configs/pixelgssgf_benchmark.yaml b/configs/pixelgssgf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8ead631c76121ac3c80ec1d617ef63cbf7608cb7 --- /dev/null +++ b/configs/pixelgssgf_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: pixelgssgf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/pup3dgs_benchmark.yaml b/configs/pup3dgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a4e224db70614b2da4cece6f05ac9710d9754c30 --- /dev/null +++ b/configs/pup3dgs_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: pup3dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/qgs_benchmark.yaml b/configs/qgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..73dc21a637a57e1d96282b42082169b63bdbc514 --- /dev/null +++ b/configs/qgs_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "qgs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/r3dg_benchmark.yaml b/configs/r3dg_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e0a4c15cd66a1b78990ce5f69a087c980a19c121 --- /dev/null +++ b/configs/r3dg_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "r3dg" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/radegs_benchmark.yaml b/configs/radegs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b47acc9d6c156e493e37ba52bd80152be640a521 --- /dev/null +++ b/configs/radegs_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "radegs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/rdogaussian_benchmark.yaml b/configs/rdogaussian_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fde42ed88de2833e2fe888e5f503e7ce0c458004 --- /dev/null +++ b/configs/rdogaussian_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: rdogaussian +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church diff --git a/configs/reactgs_benchmark.yaml b/configs/reactgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9fed97cc83a9056344a70abbda2c59ac1767d33a --- /dev/null +++ b/configs/reactgs_benchmark.yaml @@ -0,0 +1,14 @@ +method_name: reactgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/scaffoldgs_benchmark.yaml b/configs/scaffoldgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c5d94aeab60c6b4d8c8dba6ce2d114f210a032b1 --- /dev/null +++ b/configs/scaffoldgs_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: scaffoldgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/sgfgs_benchmark.yaml b/configs/sgfgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0f5f8e458ada382c5bb3a899e8707132058b66ad --- /dev/null +++ b/configs/sgfgs_benchmark.yaml @@ -0,0 +1,41 @@ +method_name: sgfgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - auditorium + - ballroom + - courtroom + - family + - francis + - horse + - lighthouse + - m60 + - museum + - palace + - panther + - playground + - temple + - train +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill diff --git a/configs/sogs_benchmark.yaml b/configs/sogs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6b0643a6fb058354a35d1378ff10dcc223b34980 --- /dev/null +++ b/configs/sogs_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "sogs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/sortfreegs_benchmark.yaml b/configs/sortfreegs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..02e1d0baa93cfedd0753952f784ab6406180ba4f --- /dev/null +++ b/configs/sortfreegs_benchmark.yaml @@ -0,0 +1,57 @@ +method_name: sortfreegs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks and Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - truck + - playground + - lighthouse + - train + - auditorium + - ballroom + - courtroom + - museum + - palace + - temple +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog + - Lego + - Materials + - Mic + - Ship +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - Playroom + - DrJohnson diff --git a/configs/specgaussian_benchmark.yaml b/configs/specgaussian_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d5fa9e629f1acdc01836365ff5eaf7080c434cb9 --- /dev/null +++ b/configs/specgaussian_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "specgaussian" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/splatfactow_benchmark.yaml b/configs/splatfactow_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7ec4abbaa9eca20311dbbd6ca46a542d34fc5efd --- /dev/null +++ b/configs/splatfactow_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "splatfactow" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/steepgs_benchmark.yaml b/configs/steepgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bc93debf34aecab08798790f9e859233bd1a43d2 --- /dev/null +++ b/configs/steepgs_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: steepgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - ignatius diff --git a/configs/sugar_benchmark.yaml b/configs/sugar_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0aada8485b0f2eee2dd179eefbb347d8913498e3 --- /dev/null +++ b/configs/sugar_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: sugar +global_settings: + iterations: 15000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church diff --git a/configs/sundae_benchmark.yaml b/configs/sundae_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..352b0e6d6c39415915838002fca3223418ac1f1a --- /dev/null +++ b/configs/sundae_benchmark.yaml @@ -0,0 +1,20 @@ +method_name: "sundae" +global_settings: + iterations: 60000 + seeds: [0] + +datasets: + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] diff --git a/configs/surfel_benchmark.yaml b/configs/surfel_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a07280490f9bb2121168338cef3a4b4b28119aee --- /dev/null +++ b/configs/surfel_benchmark.yaml @@ -0,0 +1,15 @@ +method_name: gaussian_surfel +global_settings: + iterations: 10000 + cap_gaussians: 4000000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - church + - courthouse + - ignatius + - meetingroom diff --git a/configs/tamgs_benchmark.yaml b/configs/tamgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dd63d41c01d372f413d931cbebe2e66067f22b67 --- /dev/null +++ b/configs/tamgs_benchmark.yaml @@ -0,0 +1,12 @@ +method_name: tamgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - caterpillar + - truck diff --git a/configs/texturedgs_benchmark.yaml b/configs/texturedgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..b9e68fd58623c690d6439a0c24bb2bf7cb77ff1c --- /dev/null +++ b/configs/texturedgs_benchmark.yaml @@ -0,0 +1,37 @@ +method_name: texturedgs +global_settings: + iterations: 60000 + seeds: + - 0 +datasets: +- name: Mip-NeRF 360 Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF 360 Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill +- name: Synthetic NeRF + base_path: /root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF + resolution: 1 + scenes: + - Chair + - Drums + - Ficus + - Hotdog +- name: Deep Blending + base_path: /root/autodl-tmp/dataset/deepblending_clean + resolution: 1 + scenes: + - DrJohnson + - Playroom diff --git a/configs/texturegs_benchmark.yaml b/configs/texturegs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..528958e61828c042066df0df60978bd9d1c5b127 --- /dev/null +++ b/configs/texturegs_benchmark.yaml @@ -0,0 +1,30 @@ +method_name: "texture_gs" +global_settings: + iterations: 75000 + seeds: [0] + +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/configs/topologygs_benchmark.yaml b/configs/topologygs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6b72a7da1f505ff3b513a742f2c716f81ec4f4d3 --- /dev/null +++ b/configs/topologygs_benchmark.yaml @@ -0,0 +1,15 @@ +method_name: "topologygs" +global_settings: + iterations: 30000 + seeds: [0] + +datasets: + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] diff --git a/configs/trimgs_benchmark.yaml b/configs/trimgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e5eff159bdcacaf2e2342dfb1113db27d3c2b214 --- /dev/null +++ b/configs/trimgs_benchmark.yaml @@ -0,0 +1,17 @@ +method_name: trimgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - barn + - caterpillar + - church + - courthouse + - ignatius + - meetingroom + - truck diff --git a/configs/vanilla_3dgs_ablation_benchmark.yaml b/configs/vanilla_3dgs_ablation_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1e51201a9586985af1fd6e5e1f2c503b219ed9b2 --- /dev/null +++ b/configs/vanilla_3dgs_ablation_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanilla_3dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanilla_3dgs_rest_benchmark.yaml b/configs/vanilla_3dgs_rest_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d9d034ec384177dd5680667bab3a94b4e118f01b --- /dev/null +++ b/configs/vanilla_3dgs_rest_benchmark.yaml @@ -0,0 +1,17 @@ +method_name: vanilla_3dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanilla_benchmark.yaml b/configs/vanilla_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f538970e7ac456b0d2ab80156be3db3fea511b93 --- /dev/null +++ b/configs/vanilla_benchmark.yaml @@ -0,0 +1,11 @@ +method_name: vanilla_3dgs +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - courthouse diff --git a/configs/vanilla_logged_benchmark.yaml b/configs/vanilla_logged_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3bea61a7a216297d0d7c27abb9b41a1a5ef815f5 --- /dev/null +++ b/configs/vanilla_logged_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanilla_logged +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanilla_thresh_benchmark.yaml b/configs/vanilla_thresh_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..20467c9b36a7680c0da03e6173cf97b7ec5a9373 --- /dev/null +++ b/configs/vanilla_thresh_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanilla_thresh +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_2scale_benchmark.yaml b/configs/vanillasgf_2scale_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..dfc9c019425cee3a897504f7431b734b127379d0 --- /dev/null +++ b/configs/vanillasgf_2scale_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_2scale +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_4scale_benchmark.yaml b/configs/vanillasgf_4scale_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..435bd03c7fa3744a6456ad16e539f36a057a7ed1 --- /dev/null +++ b/configs/vanillasgf_4scale_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_4scale +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_benchmark.yaml b/configs/vanillasgf_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8729803aa484aaf0cfbf757df2045e91ecd7a925 --- /dev/null +++ b/configs/vanillasgf_benchmark.yaml @@ -0,0 +1,41 @@ +method_name: vanillasgf +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - auditorium + - ballroom + - courtroom + - family + - francis + - horse + - lighthouse + - m60 + - museum + - palace + - panther + - playground + - temple + - train +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai + - counter + - kitchen + - room +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - flowers + - garden + - stump + - treehill diff --git a/configs/vanillasgf_hook_benchmark.yaml b/configs/vanillasgf_hook_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ccb3b6880b16ea73bfa4b8db070386be87056491 --- /dev/null +++ b/configs/vanillasgf_hook_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_hook +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_laplacian_benchmark.yaml b/configs/vanillasgf_laplacian_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a0d4912bb2c71e3b801ae023fab289f919a52d36 --- /dev/null +++ b/configs/vanillasgf_laplacian_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_laplacian +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_log_benchmark.yaml b/configs/vanillasgf_log_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..f719b1f4b261ccab4f196208e2effa935422c381 --- /dev/null +++ b/configs/vanillasgf_log_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_log +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_logged_benchmark.yaml b/configs/vanillasgf_logged_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..59252b684d5bcb9a4083cf5c372565e31f6d9162 --- /dev/null +++ b/configs/vanillasgf_logged_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_logged +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_nolstruct_benchmark.yaml b/configs/vanillasgf_nolstruct_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1ec0140c3d901ce7d12378cdeedf2ddba562a200 --- /dev/null +++ b/configs/vanillasgf_nolstruct_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_nolstruct +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_nomultiscale_benchmark.yaml b/configs/vanillasgf_nomultiscale_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1a7242a8d7576833bc2d4a8c90adf4eae07fac32 --- /dev/null +++ b/configs/vanillasgf_nomultiscale_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_nomultiscale +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_nosg_benchmark.yaml b/configs/vanillasgf_nosg_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c06433c067645ec3893706430662b0ceb329a5d1 --- /dev/null +++ b/configs/vanillasgf_nosg_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_nosg +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_prewitt_benchmark.yaml b/configs/vanillasgf_prewitt_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..549bee0709a4f5d6476ffb7a3c726e6f163a6378 --- /dev/null +++ b/configs/vanillasgf_prewitt_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_prewitt +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_scharr_benchmark.yaml b/configs/vanillasgf_scharr_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0612b6c56345404b1638d8495fb7f65d12c1ebab --- /dev/null +++ b/configs/vanillasgf_scharr_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_scharr +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_wgt_benchmark.yaml b/configs/vanillasgf_wgt_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..22276dcf03c00ce7f32b3f3d1756e9b657ebc625 --- /dev/null +++ b/configs/vanillasgf_wgt_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_wgt +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vanillasgf_wraw_benchmark.yaml b/configs/vanillasgf_wraw_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3d103ec1c8a31446003b632ff8eb20d77aa62f9b --- /dev/null +++ b/configs/vanillasgf_wraw_benchmark.yaml @@ -0,0 +1,24 @@ +method_name: vanillasgf_wraw +global_settings: + iterations: 30000 + seeds: + - 0 +datasets: +- name: Tanks_and_Temples + base_path: /root/autodl-tmp/dataset/tnt + resolution: 2 + scenes: + - francis + - m60 + - panther +- name: Mip-NeRF_360_Indoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 2 + scenes: + - bonsai +- name: Mip-NeRF_360_Outdoor + base_path: /root/autodl-tmp/dataset/360 + resolution: 4 + scenes: + - bicycle + - garden diff --git a/configs/vdgs_benchmark.yaml b/configs/vdgs_benchmark.yaml new file mode 100644 index 0000000000000000000000000000000000000000..841ef9f45a3ab3fc33291790c314513f2bf07cb7 --- /dev/null +++ b/configs/vdgs_benchmark.yaml @@ -0,0 +1,25 @@ +method_name: "vdgs" +global_settings: + iterations: 30000 + seeds: [0] +datasets: + - name: "Tanks and Temples" + base_path: "/root/autodl-tmp/dataset/tnt" + resolution: 2 + scenes: ["barn", "caterpillar", "truck", "playground", "lighthouse", "train", "auditorium", "ballroom", "courtroom", "museum", "palace", "temple"] + - name: "Mip-NeRF 360 Indoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 2 + scenes: ["bonsai", "counter", "kitchen", "room"] + - name: "Mip-NeRF 360 Outdoor" + base_path: "/root/autodl-tmp/dataset/360" + resolution: 4 + scenes: ["bicycle", "flowers", "garden", "stump", "treehill"] + - name: "Synthetic NeRF" + base_path: "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" + resolution: 1 + scenes: ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] + - name: "Deep Blending" + base_path: "/root/autodl-tmp/dataset/deepblending_clean" + resolution: 1 + scenes: ["Playroom", "DrJohnson"] diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/core/base_method.py b/core/base_method.py new file mode 100644 index 0000000000000000000000000000000000000000..9214ce0a50112ccb8b3bd9567f5677cdc4926432 --- /dev/null +++ b/core/base_method.py @@ -0,0 +1,69 @@ +import abc +from typing import Dict, Any, List +import torch + +class BaseMethod(abc.ABC): + """ + SplatAtlas 的统一抽象基类。 + 所有 3DGS 变体必须继承此接口，强制隔离内部逻辑与外部评测管线。 + """ + + @abc.abstractmethod + def __init__(self, dataset_config: Dict, hyperparams: Dict): + """初始化模型、优化器等，不启动训练""" + pass + + @abc.abstractmethod + def train_iteration(self, step: int) -> Dict[str, float]: + """ + 执行单步前向与反向传播，包含致密化逻辑。 + 必须返回包含指标的字典，如： + {"loss": loss, "num_gaussians": N, "pos_grad_norm": grad, "gamma_median": gamma} + """ + pass + + @abc.abstractmethod + def render(self, camera: Any) -> Dict[str, torch.Tensor]: + """无梯度渲染图像，返回 {"image": tensor, "depth": tensor(可选)}""" + pass + + @abc.abstractmethod + def save(self, save_dir: str, step: int): + """将 .ply 和 checkpoints 落盘""" + pass + + @abc.abstractmethod + def load(self, model_path: str, iteration: int): + """ + 强制实现：负责正确加载 .ply, MLP权重, 哈希表等所有必要资产。 + 绝对禁止外部脚本绕过此方法直接读取 .ply 文件。 + """ + pass + + @abc.abstractmethod + def get_spatial_centers(self) -> torch.Tensor: + """ + 核心约束：返回 [N, 3] 的拓扑中心坐标张量。 + 用于流形坍塌检测(Manifold Collapse)以及与 COLMAP SfM 点云的绝对宇宙坐标系对齐(ICP)。 + 若是隐式表面/哈希网格，须返回表面等值面的采样点。 + """ + raise NotImplementedError("该 Wrapper 尚未实现 get_spatial_centers") + + @abc.abstractmethod + def compute_physical_metrics(self, cameras: List[Any] = None) -> Dict[str, float]: + """ + 由变体自身决定如何计算表征特异性指标（如畸变度、SH能量）。 + 引入 cameras 以支持视图依赖型（View-Dependent）指标的计算。 + 若尚未实现，可抛出 NotImplementedError。 + """ + raise NotImplementedError("该 Wrapper 尚未实现 compute_physical_metrics") + + @abc.abstractmethod + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras: List[Any] = None) -> torch.Tensor: + """ + 核心重构：评估表征在指定 3D 坐标网格上的标量场（密度/不透明度/占用率）。 + - query_points: [V, 3] 的查询网格坐标张量。 + - cameras: 观察这些点的相机列表（用于打破视图依赖死锁）。 + - 返回: [V] 的标量张量，表示每个点的空间物理量。 + """ + raise NotImplementedError("该 Wrapper 尚未实现 evaluate_spatial_field") diff --git a/core/registry.py b/core/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..1296ab9dce9cab8174031b2912e3fa7cd50b1652 --- /dev/null +++ b/core/registry.py @@ -0,0 +1,10 @@ +from typing import Dict, Type +from .base_method import BaseMethod + +METHOD_REGISTRY: Dict[str, Type[BaseMethod]] = {} + +def register_method(name: str): + def decorator(cls: Type[BaseMethod]): + METHOD_REGISTRY[name] = cls + return cls + return decorator diff --git a/methods/__init__.py b/methods/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..061c6881d689d5f82f1e25b9649d1e04be352b26 --- /dev/null +++ b/methods/__init__.py @@ -0,0 +1,25 @@ +import importlib +from core.registry import METHOD_REGISTRY + +def load_method(method_name: str): + """ + 按需懒加载目标插件，实现严格的运行期隔离。 + 规避全量扫描导致的模块污染和 KeyError。 + """ + # 针对不符合 wrapper_{method} 命名规范的特例进行路由 + wrapper_router = { + "vanilla_3dgs": "wrapper_vanilla", + "tamgs": "wrapper_tamgs", + } + + module_name = wrapper_router.get(method_name, f"wrapper_{method_name}") + + try: + importlib.import_module(f".{module_name}", package=__name__) + except ImportError as e: + raise RuntimeError(f"❌ [Registry] 隔离挂载插件 {module_name} 失败，请检查其内部特异性依赖: {e}") + + if method_name not in METHOD_REGISTRY: + raise KeyError(f"❌ [Registry] 模块 {module_name} 已成功导入，但未找到方法 '{method_name}'，请检查 @register_method 装饰器。") + + return METHOD_REGISTRY[method_name] diff --git a/methods/wrapper_2dgs.py b/methods/wrapper_2dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..1c9c6c4bd45e2768ecc25364c7522a2ab8b36730 --- /dev/null +++ b/methods/wrapper_2dgs.py @@ -0,0 +1,227 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../2dgs_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("2dgs") +class TwoDGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + if "360" in str(self.args.source_path): + self.opt.lambda_dist = 100.0 + else: + self.opt.lambda_dist = 1000.0 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss_photometric = loss_target + loss_parasitic + + loss_d = render_pkg.get("distortion", torch.tensor(0.0, device="cuda")).mean() + loss_n = render_pkg.get("normal_error", torch.tensor(0.0, device="cuda")).mean() + loss = loss_photometric + self.opt.lambda_dist * loss_d + self.opt.lambda_normal * loss_n + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_2dgssgf.py b/methods/wrapper_2dgssgf.py new file mode 100644 index 0000000000000000000000000000000000000000..df29f581fde5865570ce66e0a5e9bca2bb0973da --- /dev/null +++ b/methods/wrapper_2dgssgf.py @@ -0,0 +1,292 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../2dgs_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + + +# === SGF helper functions === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("2dgssgf") +class TwoDGSSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + # 2DGS-specific: lambda_dist depends on dataset + if "360" in str(self.args.source_path): + self.opt.lambda_dist = 100.0 + else: + self.opt.lambda_dist = 1000.0 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF: structure-aware reweighted photometric loss === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + + # === 2DGS-specific: distortion + normal consistency losses === + loss_d = render_pkg.get("distortion", torch.tensor(0.0, device="cuda")).mean() + loss_n = render_pkg.get("normal_error", torch.tensor(0.0, device="cuda")).mean() + loss_2dgs_geo = self.opt.lambda_dist * loss_d + self.opt.lambda_normal * loss_n + + loss = loss_target + loss_parasitic + loss_2dgs_geo + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "loss_dist": float(loss_d), "loss_normal": float(loss_n), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted), + "iter_time_ms": float((time.perf_counter() - _iter_start) * 1000), + "vram_allocated_GB": float(torch.cuda.memory_allocated() / 1024**3), + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_3dgs_dr.py b/methods/wrapper_3dgs_dr.py new file mode 100644 index 0000000000000000000000000000000000000000..6c5c99898a89d6b0531401c90d87aebdc4a6f6e6 --- /dev/null +++ b/methods/wrapper_3dgs_dr.py @@ -0,0 +1,258 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3DGS-DR'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("3dgs_dr") +class ThreeDGSDRWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + + if step > self.opt.feature_rest_from_iter and step % 1000 == 0: + self.gaussians.oneupSHdegree() + + initial_stage = step <= self.opt.init_until_iter + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, initial_stage=initial_stage) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + if self.opt.use_env_scope and 'refl_strength_map' in render_pkg: + center = torch.tensor([float(c) for c in self.opt.env_scope_center], device='cuda') + outside_msk = torch.sum((self.gaussians.get_xyz - center[None])**2, dim=-1) > self.opt.env_scope_radius**2 + refls = self.gaussians.get_refl + refl_msk_loss = refls[outside_msk].mean() + loss += 0.4 * refl_msk_loss + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + DENSIFY_UNTIL_ITER = self.opt.densify_until_iter + self.opt.longer_prop_iter + NORMAL_PROP_UNTIL_ITER = self.opt.normal_prop_until_iter + self.opt.longer_prop_iter + + if step < DENSIFY_UNTIL_ITER: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step <= self.opt.init_until_iter: + opacity_reset_intval = 3000 + densification_interval = 100 + elif step <= NORMAL_PROP_UNTIL_ITER: + opacity_reset_intval = 3000 + densification_interval = self.opt.densification_interval_when_prop + else: + opacity_reset_intval = 3000 + densification_interval = 100 + + if step > self.opt.densify_from_iter and step % densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.prune_opacity_threshold, self.scene.cameras_extent, size_threshold) + + HAS_RESET0 = False + if step % opacity_reset_intval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + HAS_RESET0 = True + self.gaussians.reset_opacity0() + self.gaussians.reset_refl() + + if self.opt.opac_lr0_interval > 0 and (self.opt.init_until_iter < step <= NORMAL_PROP_UNTIL_ITER) and step % self.opt.opac_lr0_interval == 0: + self.gaussians.set_opacity_lr(self.opt.opacity_lr) + + if (self.opt.init_until_iter < step <= NORMAL_PROP_UNTIL_ITER) and step % 1000 == 0: + if not HAS_RESET0: + self.gaussians.reset_opacity1() + self.gaussians.dist_color() + self.gaussians.reset_scale() + if self.opt.opac_lr0_interval > 0 and step != NORMAL_PROP_UNTIL_ITER: + self.gaussians.set_opacity_lr(0.0) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + histograms["refl_strength"] = torch.sigmoid(self.gaussians._refl_strength).clone().detach() + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, initial_stage=False) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + metrics["refl_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._refl_strength))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_3dgsmcmc.py b/methods/wrapper_3dgsmcmc.py new file mode 100644 index 0000000000000000000000000000000000000000..00431aff75e85f4707400f3bc30c05fec8be553d --- /dev/null +++ b/methods/wrapper_3dgsmcmc.py @@ -0,0 +1,301 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsmcmc_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.gaussian_model import build_scaling_rotation +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("3dgsmcmc") +class ThreeDGSMCMCWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + scene_name = os.path.basename(os.path.normpath(self.args.source_path)).lower() + indoor_scenes = ["bonsai", "counter", "kitchen", "room", "playroom", "drjohnson", "train"] + is_indoor = "indoor" in self.args.source_path.lower() or scene_name in indoor_scenes + self.args.cap_max = 2000000 if is_indoor else 4000000 + + self.dataset = self.lp.extract(self.args) + self.dataset.cap_max = self.args.cap_max + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + xyz_lr = self.gaussians.update_learning_rate(step) + + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + loss_parasitic = self.opt.opacity_reg * torch.abs(self.gaussians.get_opacity).mean() + self.opt.scale_reg * torch.abs(self.gaussians.get_scaling).mean() + loss = loss_target + loss_parasitic + + sti_metrics = { + "sti_spatial": 0.0, + "sti_geometry": 0.0, + "sti_opacity": 0.0, + "sti_appearance": 0.0 + } + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_t = {} + for group in self.gaussians.optimizer.param_groups: + name = group["name"] + p = group['params'][0] + if p.grad is not None: + grads_t[name] = p.grad.clone() + else: + grads_t[name] = torch.zeros_like(p) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_p = {} + states_v = {} + lrs = {} + epsilons = {} + for group in self.gaussians.optimizer.param_groups: + name = group["name"] + p = group['params'][0] + if p.grad is not None: + grads_p[name] = p.grad.clone() + else: + grads_p[name] = torch.zeros_like(p) + + state = self.gaussians.optimizer.state.get(p, {}) + states_v[name] = state.get("exp_avg_sq", torch.zeros_like(p)) + lrs[name] = group["lr"] + epsilons[name] = group.get("eps", 1e-15) + + semantic_groups = { + "spatial": ["xyz"], + "geometry": ["scaling", "rotation"], + "opacity": ["opacity"], + "appearance": ["f_dc", "f_rest"] + } + + for sg_name, param_names in semantic_groups.items(): + u_t_list = [] + u_p_list = [] + for pname in param_names: + if pname in grads_t: + gt = grads_t[pname] + gp = grads_p[pname] + vt = states_v[pname] + lr = lrs[pname] + eps = epsilons[pname] + + denom = torch.sqrt(vt) + eps + ut = torch.nan_to_num((lr / denom) * gt, nan=0.0, posinf=0.0, neginf=0.0) + up = torch.nan_to_num((lr / denom) * gp, nan=0.0, posinf=0.0, neginf=0.0) + + N = ut.shape[0] + u_t_list.append(ut.view(N, -1)) + u_p_list.append(up.view(N, -1)) + + if len(u_t_list) > 0: + M_UT = torch.cat(u_t_list, dim=1) + M_UP = torch.cat(u_p_list, dim=1) + + norm_ut = torch.norm(M_UT, dim=1) + norm_up = torch.norm(M_UP, dim=1) + + valid_mask = (norm_ut > 0) & (norm_up > 0) + + if valid_mask.any(): + cos_sim = F.cosine_similarity(M_UT[valid_mask], M_UP[valid_mask], dim=1) + share = norm_up[valid_mask] / (norm_ut[valid_mask] + norm_up[valid_mask] + 1e-7) + ti = share * torch.clamp(-cos_sim, min=0.0) + sti_metrics[f"sti_{sg_name}"] = float(ti.mean().item()) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if self.opt.densify_from_iter < step < self.opt.densify_until_iter and step % self.opt.densification_interval == 0: + dead_mask = (self.gaussians.get_opacity <= 0.005).squeeze(-1) + self.gaussians.relocate_gs(dead_mask=dead_mask) + self.gaussians.add_new_gs(cap_max=self.dataset.cap_max) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + L = build_scaling_rotation(self.gaussians.get_scaling, self.gaussians.get_rotation) + actual_covariance = L @ L.transpose(1, 2) + + def op_sigmoid(x, k=100, x0=0.995): + return 1.0 / (1.0 + torch.exp(-k * (x - x0))) + + noise = torch.randn_like(self.gaussians._xyz) * op_sigmoid(1.0 - self.gaussians.get_opacity) * self.opt.noise_lr * xyz_lr + noise = torch.bmm(actual_covariance, noise.unsqueeze(-1)).squeeze(-1) + self.gaussians._xyz.add_(noise) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) + } + metrics.update(sti_metrics) + + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + opacity = self.gaussians.get_opacity.clone().detach() + scales = self.gaussians.get_scaling.clone().detach() + + histograms["opacity"] = opacity + histograms["scaling"] = scales + + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + mcmc_modulation = 1.0 / (1.0 + torch.exp(-100.0 * ((1.0 - opacity) - 0.995))) + histograms["mcmc_noise_modulation"] = mcmc_modulation.squeeze(-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + bg = torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device="cuda") + render_pkg = native_render(camera, self.gaussians, self.pipe, bg) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.gaussians.get_opacity)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2.0*(x*z + w*y), 2.0*(y*z - w*x), 1.0-2.0*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, self.gaussians.get_opacity.squeeze(-1) + scales = self.gaussians.get_scaling + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze(-1).pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_3dgut.py b/methods/wrapper_3dgut.py new file mode 100644 index 0000000000000000000000000000000000000000..0c9d08abfbc9dfd26e8a7a9be3eabbbdd20db207 --- /dev/null +++ b/methods/wrapper_3dgut.py @@ -0,0 +1,258 @@ +import os +import sys +import torch +import numpy as np +import torch.nn.functional as F +from omegaconf import OmegaConf +from hydra import initialize_config_dir, compose +from hydra.core.global_hydra import GlobalHydra + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgrut'))) +from threedgrut.trainer import Trainer3DGRUT +from threedgrut.optimizers import SelectiveAdam + +class MockScene: + def getTrainCameras(self): + return [] + def getTestCameras(self): + return [] + +@register_method("3dgut") +class ThreeDGUTWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.track_decoupling = hyperparams.get("track_decoupling", False) + + GlobalHydra.instance().clear() + config_dir = os.path.abspath("/root/autodl-tmp/3dgrut/configs") + + with initialize_config_dir(version_base=None, config_dir=config_dir): + ds_colmap = "dataset=colmap" + try: + compose(config_name="base_gs", overrides=["dataset=colmap"]) + except Exception: + ds_colmap = "+dataset=colmap" + + ds_nerf = "dataset=nerf" + try: + compose(config_name="base_gs", overrides=["dataset=nerf"]) + except Exception: + ds_nerf = "+dataset=nerf" + + overrides = [ + "+render=3dgut", + f"++dataset.path={dataset_config['source_path']}", + f"++path={dataset_config['source_path']}", + f"++out_dir={dataset_config['model_path']}", + "++experiment_name=benchmark", + "++n_iterations=30000", + "++loss.use_l1=True", + "++loss.lambda_l1=0.8", + "++loss.use_ssim=True", + "++loss.lambda_ssim=0.2", + "++with_gui=False", + "++with_viser_gui=False", + ] + + if "Synthetic_NeRF" in str(dataset_config["source_path"]): + overrides.append(ds_nerf) + else: + overrides.append(ds_colmap) + overrides.append(f"++dataset.downsample_factor={dataset_config.get('resolution', 1)}") + + conf = compose(config_name="base_gs", overrides=overrides) + + OmegaConf.set_struct(conf, False) + + if "initialization" not in conf: + conf.initialization = OmegaConf.create({}) + + if "Synthetic_NeRF" in str(dataset_config["source_path"]): + conf.initialization.method = "random" + conf.initialization.num_gaussians = 100000 + else: + conf.initialization.method = "colmap" + + conf.initialization.use_observation_points = False + + if not hasattr(conf.dataset, "test_split_interval"): + conf.dataset.test_split_interval = 8 + if not hasattr(conf.dataset, "eval"): + conf.dataset.eval = True + + self.trainer = Trainer3DGRUT(conf) + self.model = self.trainer.model + self.data_iter = iter(self.trainer.train_dataloader) + self.last_n_gaussians = self.model.num_gaussians + self.scene = MockScene() + + def train_iteration(self, step): + try: + batch = next(self.data_iter) + except StopIteration: + self.data_iter = iter(self.trainer.train_dataloader) + batch = next(self.data_iter) + + gpu_batch = self.trainer.train_dataset.get_gpu_batch_with_intrinsics(batch) + outputs = self.model(gpu_batch, train=True, frame_id=step) + batch_losses = self.trainer.get_losses(gpu_batch, outputs) + + loss_target = batch_losses.get("l1_loss", torch.tensor(0.0, device="cuda")) + loss_parasitic = batch_losses.get("ssim_loss", torch.tensor(0.0, device="cuda")) + loss = batch_losses.get("total_loss", loss_target + loss_parasitic) + + self.trainer.strategy.pre_backward( + step=step, + scene_extent=self.trainer.scene_extent, + train_dataset=self.trainer.train_dataset, + batch=gpu_batch, + writer=None, + ) + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.model.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.model.positions.grad.clone() if self.model.positions.grad is not None else torch.zeros_like(self.model.positions) + + self.model.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.model.positions.grad.clone() if self.model.positions.grad is not None else torch.zeros_like(self.model.positions) + + state = self.model.optimizer.state.get(self.model.positions, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = self.model.optimizer.param_groups[0]["lr"] + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + else: + u_t = grad_target + u_p = grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.model.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + self.trainer.strategy.post_backward( + step=step, + scene_extent=self.trainer.scene_extent, + train_dataset=self.trainer.train_dataset, + batch=gpu_batch, + writer=None, + ) + + if "mog_visibility" in outputs and isinstance(self.model.optimizer, SelectiveAdam): + self.model.optimizer.step(outputs["mog_visibility"]) + else: + self.model.optimizer.step() + + self.model.optimizer.zero_grad() + self.model.scheduler_step(step) + + self.trainer.strategy.post_optimizer_step( + step=step, + scene_extent=self.trainer.scene_extent, + train_dataset=self.trainer.train_dataset, + batch=gpu_batch, + writer=None, + ) + + num_gaussians = self.model.num_gaussians + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "hits_mean": float(outputs.get("hits_count", torch.tensor(0.0)).float().mean()) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = self.model.get_density().clone().detach() + scales = self.model.get_scale().clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.model.features_albedo.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + gpu_batch = self.trainer.val_dataset.get_gpu_batch_with_intrinsics(camera) + outputs = self.model(gpu_batch, train=False) + return {"image": outputs["pred_rgb"], "depth": outputs.get("pred_dist", None)} + + def save(self, save_dir, step): + self.trainer.save_checkpoint(last_checkpoint=(step >= 30000)) + + def load(self, model_path, iteration): + ckpt_path = os.path.join(model_path, f"ours_{iteration}", f"ckpt_{iteration}.pt") + if not os.path.exists(ckpt_path): + ckpt_path = os.path.join(model_path, "ckpt_last.pt") + checkpoint = torch.load(ckpt_path, map_location="cuda", weights_only=False) + self.model.init_from_checkpoint(checkpoint, setup_optimizer=False) + + def get_spatial_centers(self): + return self.model.positions + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.model.get_scale() + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.model.get_density())) + dc = self.model.features_albedo + rest = self.model.features_specular + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.model.rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.model.positions + opacities = self.model.get_density().squeeze() + scales = self.model.get_scale() + sigma_sq = scales[:, :2].max(dim=1)[0].pow(2) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_aaags.py b/methods/wrapper_aaags.py new file mode 100644 index 0000000000000000000000000000000000000000..ee5bd6af2741091c5d9420b18d2c3356737bc733 --- /dev/null +++ b/methods/wrapper_aaags.py @@ -0,0 +1,283 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/AAA-Gaussians_offy') +from utils.loss_utils import l1_loss +from utils.loss_utils import ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.gaussian_model import build_scaling_rotation +from arguments import ModelParams, PipelineParams, OptimizationParams, SplattingSettings + +@register_method("aaags") +class AAAGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.ss = SplattingSettings(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.splat_args = self.ss.get_settings(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_effective_step(self, params): + steps = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get("exp_avg_sq", torch.zeros_like(p.grad)) + lr = 0.0 + for pg in self.gaussians.optimizer.param_groups: + if any(p is param for param in pg['params']): + lr = pg['lr'] + break + step = (lr / (torch.sqrt(v) + 1e-15)) * p.grad.clone() + steps.append(step.reshape(-1)) + if steps: + return torch.cat(steps) + return None + + def train_iteration(self, step): + xyz_lr = self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, splat_args=self.splat_args) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + \ + self.opt.opacity_reg * torch.abs(self.gaussians.get_opacity).mean() + \ + self.opt.scale_reg * torch.abs(self.gaussians.get_scaling).mean() + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest] + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + steps_target = {} + for group_name, params in param_groups_map.items(): + steps_target[group_name] = self._get_effective_step(params) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + steps_parasitic = {} + for group_name, params in param_groups_map.items(): + steps_parasitic[group_name] = self._get_effective_step(params) + + for group_name in param_groups_map: + u_t = steps_target.get(group_name) + u_p = steps_parasitic.get(group_name) + if u_t is not None and u_p is not None and u_t.norm() > 0 and u_p.norm() > 0: + cos = float(F.cosine_similarity(u_t.unsqueeze(0), u_p.unsqueeze(0))) + r = float(u_p.norm() / (u_t.norm() + u_p.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + dead_mask = (self.gaussians.get_opacity <= 0.005).squeeze(-1) + self.gaussians.relocate_gs(viewpoint_cam.camera_center, dead_mask=dead_mask) + self.gaussians.add_new_gs(cap_max=self.dataset.cap_max) + + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + L = build_scaling_rotation(self.gaussians.get_scaling, self.gaussians.get_rotation) + actual_covariance = L @ L.transpose(1, 2) + + def op_sigmoid(x, k=100, x0=0.995): + return 1 / (1 + torch.exp(-k * (x - x0))) + + noise = torch.randn_like(self.gaussians._xyz) * (op_sigmoid(1 - self.gaussians.get_opacity)) * self.opt.noise_lr * xyz_lr + noise = torch.bmm(actual_covariance, noise.unsqueeze(-1)).squeeze(-1) + self.gaussians._xyz.add_(noise) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, splat_args=self.splat_args) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_absgs.py b/methods/wrapper_absgs.py new file mode 100644 index 0000000000000000000000000000000000000000..3ad615dda8ea70be348411c37fdb326599e0d7ca --- /dev/null +++ b/methods/wrapper_absgs.py @@ -0,0 +1,285 @@ +import os +import sys +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../absgs_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("absgs") +class AbsGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.percent_dense = 0.001 + + if "360" in str(self.args.source_path): + self.opt.lambda_dist = 100.0 + else: + self.opt.lambda_dist = 1000.0 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gs_w = render_pkg.get("gs_w", torch.zeros_like(radii)) + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + target_updates = [] + for group in self.gaussians.optimizer.param_groups: + p = group['params'][0] + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and 'exp_avg_sq' in state: + v = state['exp_avg_sq'] + lr = group['lr'] + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + target_updates.append(u.view(u.shape[0], -1)) + else: + target_updates.append(torch.zeros_like(p).view(p.shape[0], -1)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + parasitic_updates = [] + for group in self.gaussians.optimizer.param_groups: + p = group['params'][0] + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and 'exp_avg_sq' in state: + v = state['exp_avg_sq'] + lr = group['lr'] + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + parasitic_updates.append(u.view(u.shape[0], -1)) + else: + parasitic_updates.append(torch.zeros_like(p).view(p.shape[0], -1)) + + u_t = torch.cat(target_updates, dim=1) + u_p = torch.cat(parasitic_updates, dim=1) + + norm_t = torch.norm(u_t, dim=1) + norm_p = torch.norm(u_p, dim=1) + + valid_mask = (norm_t > 0) & (norm_p > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float((norm_p / (norm_t + norm_p + 1e-7)).mean()) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.gaussians.max_weight[visibility_filter] = torch.max(self.gaussians.max_weight[visibility_filter], gs_w[visibility_filter]) + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.densify_grad_abs_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reduce_interval == 0 and self.opt.use_reduce: + self.gaussians.reduce_opacity() + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step > self.opt.densify_from_iter and step < self.opt.prune_until_iter and self.opt.use_prune_weight: + img_num = len(self.viewpoint_stack) if len(self.viewpoint_stack) > 0 else 1 + if step % img_num == 0 and step % self.opt.opacity_reset_interval > img_num: + prune_mask = (self.gaussians.max_weight < self.opt.min_weight).squeeze() + self.gaussians.prune_points(prune_mask) + self.gaussians.max_weight *= 0 + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + denom_mask = self.gaussians.denom > 0 + grad_reg_mean = float((self.gaussians.xyz_gradient_accum[denom_mask] / self.gaussians.denom[denom_mask]).mean()) if denom_mask.any() else 0.0 + grad_abs_mean = float((self.gaussians.xyz_gradient_accum_abs[denom_mask] / self.gaussians.denom[denom_mask]).mean()) if denom_mask.any() else 0.0 + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "grad_regular_mean": float(grad_reg_mean), "grad_abs_mean": float(grad_abs_mean), + "max_weight_mean": float(self.gaussians.max_weight.mean()) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_absgssgf.py b/methods/wrapper_absgssgf.py new file mode 100644 index 0000000000000000000000000000000000000000..9a72c8b7be3befca0e463ea96fd09b60206d9e36 --- /dev/null +++ b/methods/wrapper_absgssgf.py @@ -0,0 +1,336 @@ +import os +import sys +import math +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../absgs_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + +@register_method("absgssgf") +class AbsGSSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.percent_dense = 0.001 + + if "360" in str(self.args.source_path): + self.opt.lambda_dist = 100.0 + else: + self.opt.lambda_dist = 1000.0 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gs_w = render_pkg.get("gs_w", torch.zeros_like(radii)) + + gt_image = viewpoint_cam.original_image.cuda() + + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_struct + loss_parasitic = loss_ssim_reweighted + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + target_updates = [] + for group in self.gaussians.optimizer.param_groups: + p = group['params'][0] + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and 'exp_avg_sq' in state: + v = state['exp_avg_sq'] + lr = group['lr'] + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + target_updates.append(u.view(u.shape[0], -1)) + else: + target_updates.append(torch.zeros_like(p).view(p.shape[0], -1)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + parasitic_updates = [] + for group in self.gaussians.optimizer.param_groups: + p = group['params'][0] + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and 'exp_avg_sq' in state: + v = state['exp_avg_sq'] + lr = group['lr'] + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + parasitic_updates.append(u.view(u.shape[0], -1)) + else: + parasitic_updates.append(torch.zeros_like(p).view(p.shape[0], -1)) + + u_t = torch.cat(target_updates, dim=1) + u_p = torch.cat(parasitic_updates, dim=1) + + norm_t = torch.norm(u_t, dim=1) + norm_p = torch.norm(u_p, dim=1) + + valid_mask = (norm_t > 0) & (norm_p > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float((norm_p / (norm_t + norm_p + 1e-7)).mean()) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.gaussians.max_weight[visibility_filter] = torch.max(self.gaussians.max_weight[visibility_filter], gs_w[visibility_filter]) + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.densify_grad_abs_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reduce_interval == 0 and self.opt.use_reduce: + self.gaussians.reduce_opacity() + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step > self.opt.densify_from_iter and step < self.opt.prune_until_iter and self.opt.use_prune_weight: + img_num = len(self.viewpoint_stack) if len(self.viewpoint_stack) > 0 else 1 + if step % img_num == 0 and step % self.opt.opacity_reset_interval > img_num: + prune_mask = (self.gaussians.max_weight < self.opt.min_weight).squeeze() + self.gaussians.prune_points(prune_mask) + self.gaussians.max_weight *= 0 + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + denom_mask = self.gaussians.denom > 0 + grad_reg_mean = float((self.gaussians.xyz_gradient_accum[denom_mask] / self.gaussians.denom[denom_mask]).mean()) if denom_mask.any() else 0.0 + grad_abs_mean = float((self.gaussians.xyz_gradient_accum_abs[denom_mask] / self.gaussians.denom[denom_mask]).mean()) if denom_mask.any() else 0.0 + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "grad_regular_mean": float(grad_reg_mean), "grad_abs_mean": float(grad_abs_mean), + "max_weight_mean": float(self.gaussians.max_weight.mean()), + "sgf_w_mean": float(W.mean()), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_analyticsplatting.py b/methods/wrapper_analyticsplatting.py new file mode 100644 index 0000000000000000000000000000000000000000..04a558714714b9415f0841fbb8b8b9390eed86c4 --- /dev/null +++ b/methods/wrapper_analyticsplatting.py @@ -0,0 +1,300 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +_AS_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Analytic-Splatting')) +assert os.path.isdir(_AS_ROOT), f"Analytic-Splatting repo not found at {_AS_ROOT}" +# insert(0) so its `scene` / `arguments` modules win over the generic 3DGS ones in main_train sys.path +if _AS_ROOT not in sys.path: + sys.path.insert(0, _AS_ROOT) +# evict generic 3DGS modules that may already be cached, so analytic-splatting's versions get loaded +for _m in list(sys.modules): + if _m == 'scene' or _m.startswith('scene.') \ + or _m == 'arguments' or _m.startswith('arguments.') \ + or _m == 'gaussian_renderer' or _m.startswith('gaussian_renderer.') \ + or _m == 'utils' or _m.startswith('utils.'): + sys.modules.pop(_m, None) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("analyticsplatting") +class AnalyticSplattingWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.parser.add_argument("--filter3d", action="store_true") + self.parser.add_argument("--dense", action="store_true") + + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.filter3d = True + + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + + self.scene_scales = [1.0, 2.0, 4.0, 8.0] + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians, resolution_scales=self.scene_scales) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + + self.trainCameras = [] + for scale in self.scene_scales: + self.trainCameras += self.scene.getTrainCameras(scale).copy() + + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + self.viewpoint_stack = self.trainCameras.copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_effective_step(self, param, grad, group_lr): + state = self.gaussians.optimizer.state.get(param) + v = state.get("exp_avg_sq") if state is not None else torch.ones_like(param) * 1e-8 + return (group_lr / (torch.sqrt(v) + 1e-15)) * grad + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.trainCameras.copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, filter3d=self.args.filter3d) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target_full = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + grads_target_dict = {p: p.grad.clone() if p.grad is not None else torch.zeros_like(p) for group in self.gaussians.optimizer.param_groups for p in group["params"]} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic_full = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + grads_parasitic_dict = {p: p.grad.clone() if p.grad is not None else torch.zeros_like(p) for group in self.gaussians.optimizer.param_groups for p in group["params"]} + + valid_mask = (torch.norm(grad_target_full, dim=1) > 0) & (torch.norm(grad_parasitic_full, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target_full[valid_mask], grad_parasitic_full[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic_full, dim=1).mean() / (torch.norm(grad_target_full, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest] + } + + for group_name, params in param_groups_map.items(): + U_T_list, U_P_list = [], [] + for p in params: + group_lr = 0.0 + for opt_g in self.gaussians.optimizer.param_groups: + if any(p is q for q in opt_g["params"]): + group_lr = opt_g["lr"] + break + + ut = self._get_effective_step(p, grads_target_dict[p], group_lr) + up = self._get_effective_step(p, grads_parasitic_dict[p], group_lr) + U_T_list.append(ut.reshape(-1)) + U_P_list.append(up.reshape(-1)) + + U_T = torch.cat(U_T_list) + U_P = torch.cat(U_P_list) + + if U_T.norm() > 0 and U_P.norm() > 0: + cos = float(F.cosine_similarity(U_T.unsqueeze(0), U_P.unsqueeze(0))) + r = float(U_P.norm() / (U_T.norm() + U_P.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, dense=self.args.dense) + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter: + if step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + if hasattr(self.gaussians, "filter_3D"): + histograms["filter_3D"] = self.gaussians.filter_3D.clone().detach() + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, filter3d=self.args.filter3d) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + if hasattr(self.gaussians, "filter_3D"): + metrics["filter_3d_mean"] = float(torch.mean(self.gaussians.filter_3D)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + if hasattr(self.gaussians, "filter_3D"): + filter_sq = self.gaussians.filter_3D.squeeze().pow(2) + sigma_sq = sigma_sq + filter_sq + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_atomgs.py b/methods/wrapper_atomgs.py new file mode 100644 index 0000000000000000000000000000000000000000..f3f39caeb2d312b15361bcb0e538f883388fca50 --- /dev/null +++ b/methods/wrapper_atomgs.py @@ -0,0 +1,262 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../AtomGS_official'))) +from utils.loss_utils import l1_loss, edge_aware_normal_loss, ms_ssim +from utils.image_utils import depth_to_normal +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("atomgs") +class AtomGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + L_ms_ssim = 1.0 - ms_ssim(image, gt_image) + loss_photometric = (1.0 - self.opt.lambda_ms_ssim) * Ll1 + self.opt.lambda_ms_ssim * L_ms_ssim + loss = loss_photometric + + Lnormal = torch.tensor(0.0, device="cuda") + if step < self.opt.atom_proliferation_until: + Lnormal = edge_aware_normal_loss(gt_image, depth_to_normal(render_pkg["mean_depth"], viewpoint_cam).permute(2, 0, 1)) + loss = loss + self.opt.lambda_normal * Lnormal + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and step < self.opt.atom_proliferation_until: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_photometric.backward(retain_graph=True) + g_T = {g['name']: g['params'][0].grad.clone() if g['params'][0].grad is not None else torch.zeros_like(g['params'][0]) for g in self.gaussians.optimizer.param_groups} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic = self.opt.lambda_normal * Lnormal + loss_parasitic.backward(retain_graph=True) + g_P = {g['name']: g['params'][0].grad.clone() if g['params'][0].grad is not None else torch.zeros_like(g['params'][0]) for g in self.gaussians.optimizer.param_groups} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + + groups_map = { + "spatial": ["xyz"], + "geometry": ["scaling", "rotation"], + "opacity": ["opacity"], + "appearance": ["f_dc", "f_rest"] + } + + sim_sum = 0.0 + ratio_sum = 0.0 + total_active = 0 + + for g_name, keys in groups_map.items(): + U_T_g = [] + U_P_g = [] + for k in keys: + for group in self.gaussians.optimizer.param_groups: + if group['name'] == k: + p = group['params'][0] + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and 'exp_avg_sq' in state: + v = state['exp_avg_sq'] + lr = group['lr'] + U_T_g.append((lr / (torch.sqrt(v) + 1e-15)) * g_T[k]) + U_P_g.append((lr / (torch.sqrt(v) + 1e-15)) * g_P[k]) + + if len(U_T_g) > 0: + U_T_vec = torch.cat([x.flatten() for x in U_T_g]) + U_P_vec = torch.cat([x.flatten() for x in U_P_g]) + + norm_P = torch.norm(U_P_vec) + if norm_P > 1e-7: + norm_T = torch.norm(U_T_vec) + cos_sim = float(F.cosine_similarity(U_T_vec.unsqueeze(0), U_P_vec.unsqueeze(0)).item()) + ratio = float(norm_P / (norm_T + norm_P + 1e-7)) + sim_sum += cos_sim + ratio_sum += ratio + total_active += 1 + + if total_active > 0: + grad_cos_sim = sim_sum / total_active + parasitic_ratio = ratio_sum / total_active + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + self.gaussians.densify_and_prune(self.opt.clone_threshold, min(self.opt.split_threshold * step / self.opt.warm_up_until, self.opt.split_threshold), self.opt.prune_threshold) + if step % self.opt.densification_interval == 0 and step < self.opt.atom_proliferation_until: + self.gaussians.atomize() + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(Ll1), + "loss_ssim": float(L_ms_ssim), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "atom_scale": float(getattr(self.gaussians, "atom_scale", 0.0)) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("mean_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_bags.py b/methods/wrapper_bags.py new file mode 100644 index 0000000000000000000000000000000000000000..3df3047ab740bc7c153f214d8e80e8141be4f197 --- /dev/null +++ b/methods/wrapper_bags.py @@ -0,0 +1,377 @@ +import os +import sys +import math +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/BAGS_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def tv_loss(grids): + number_of_grids = grids.shape[0] + h_tv_count = grids[:, :, 1:, :].shape[1] * grids[:, :, 1:, :].shape[2] * grids[:, :, 1:, :].shape[3] + w_tv_count = grids[:, :, :, 1:].shape[1] * grids[:, :, :, 1:].shape[2] * grids[:, :, :, 1:].shape[3] + h_tv = torch.pow((grids[:, :, 1:, :] - grids[:, :, :-1, :]), 2).sum() + w_tv = torch.pow((grids[:, :, :, 1:] - grids[:, :, :, :-1]), 2).sum() + return 2 * (h_tv / h_tv_count + w_tv / w_tv_count) / number_of_grids + +def get_emb(sin_inp): + emb = torch.stack((sin_inp.sin(), sin_inp.cos()), dim=-1) + return torch.flatten(emb, -2, -1) + +def get_2d_emb(batch_size, x, y, out_ch, device): + out_ch = int(np.ceil(out_ch / 4) * 2) + inv_freq = 1.0 / (10000 ** (torch.arange(0, out_ch, 2).float() / out_ch)) + pos_x = torch.arange(x, device=device).type(inv_freq.type()) * 2 * np.pi / x + pos_y = torch.arange(y, device=device).type(inv_freq.type()) * 2 * np.pi / y + sin_inp_x = torch.einsum("i,j->ij", pos_x, inv_freq) + sin_inp_y = torch.einsum("i,j->ij", pos_y, inv_freq) + emb_x = get_emb(sin_inp_x).unsqueeze(1) + emb_y = get_emb(sin_inp_y) + emb = torch.zeros((x, y, out_ch * 2), device=device) + emb[:, :, : out_ch] = emb_x + emb[:, :, out_ch : 2 * out_ch] = emb_y + return emb[None, :, :, :].repeat(batch_size, 1, 1, 1) + +@register_method("bags") +class BAGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.iterations = 60000 + self.opt.ms_steps = 6000 + self.upsample_iter = [3000, 6000] + + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset) + + inp_shape = [len(self.scene.getTrainCameras()), int(np.round(self.scene.orig_h/self.args.resolution)), int(np.round(self.scene.orig_w/self.args.resolution))] + + self.gaussians = GaussianModel( + self.dataset.sh_degree, inp_shape, + ks1=self.dataset.kernel_size1, ks2=self.dataset.kernel_size2, + ks3=self.dataset.kernel_size3, ks_ss=self.dataset.kernel_size_ss, + not_use_rgbd=self.opt.not_use_rgbd, not_use_pe=self.opt.not_use_pe + ) + + self.scene.load_gaussian(self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + + self.trainCameras_scale4 = self.scene.getTrainCameras(scale=4.0).copy() + self.trainCameras_scale2 = self.scene.getTrainCameras(scale=2.0).copy() + self.trainCameras_scale1 = self.scene.getTrainCameras(scale=1.0).copy() + + self.viewpoint_stack = self.trainCameras_scale4.copy() + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale4) + + self.unfold1 = torch.nn.Unfold(kernel_size=(self.dataset.kernel_size1, self.dataset.kernel_size1), padding=self.dataset.kernel_size1 // 2).cuda() + self.unfold2 = torch.nn.Unfold(kernel_size=(self.dataset.kernel_size2, self.dataset.kernel_size2), padding=self.dataset.kernel_size2 // 2).cuda() + self.unfold3 = torch.nn.Unfold(kernel_size=(self.dataset.kernel_size3, self.dataset.kernel_size3), padding=self.dataset.kernel_size3 // 2).cuda() + if self.dataset.kernel_size_ss != self.dataset.kernel_size3: + self.opt.use_another_mlp = True + self.unfold_ss = torch.nn.Unfold(kernel_size=(self.dataset.kernel_size_ss, self.dataset.kernel_size_ss), padding=self.dataset.kernel_size_ss // 2).cuda() + else: + self.unfold_ss = self.unfold3 + + self.last_n_gaussians = self.gaussians.get_xyz.shape[0] + + def train_iteration(self, step): + ori_iter = step + adj_iter = step + + if ori_iter == self.upsample_iter[0]: + self.viewpoint_stack = self.trainCameras_scale2.copy() + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale2) + elif ori_iter == self.upsample_iter[1]: + self.viewpoint_stack = self.trainCameras_scale1.copy() + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale1) + + if ori_iter > self.opt.ms_steps: + adj_iter = step - self.opt.ms_steps + + self.gaussians.update_learning_rate(adj_iter) + if adj_iter % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + if ori_iter >= self.upsample_iter[0] and ori_iter < self.upsample_iter[1]: + self.viewpoint_stack = self.trainCameras_scale2.copy() + elif ori_iter >= self.upsample_iter[1]: + self.viewpoint_stack = self.trainCameras_scale1.copy() + else: + self.viewpoint_stack = self.trainCameras_scale4.copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + subpixel_offset = None + if self.dataset.ray_jitter: + subpixel_offset = torch.rand((int(viewpoint_cam.image_height), int(viewpoint_cam.image_width), 2), dtype=torch.float32, device="cuda") - 0.5 + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size, subpixel_offset=subpixel_offset) + image, depth, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["depth"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + loss_target = torch.tensor(0.0, device="cuda") + loss_parasitic = torch.tensor(0.0, device="cuda") + + if adj_iter > 250: + shuffle_rgb = image.unsqueeze(0) + shuffle_depth = depth.unsqueeze(0) - depth.min() + shuffle_depth = shuffle_depth / (shuffle_depth.max() + 1e-7) + pos_enc = get_2d_emb(1, shuffle_rgb.shape[-2], shuffle_rgb.shape[-1], 16, torch.device("cuda")) + + if ori_iter < 3000: + kernel_weights, mask = self.gaussians.mlp_rgb_ms(0, pos_enc, torch.cat([shuffle_rgb, shuffle_depth], 1).detach(), ori_iter) + patches = self.unfold1(shuffle_rgb).view(1, 3, self.dataset.kernel_size1 ** 2, shuffle_rgb.shape[-2], shuffle_rgb.shape[-1]) + elif ori_iter >= 3000 and ori_iter < 6000: + kernel_weights, mask = self.gaussians.mlp_rgb_ms(0, pos_enc, torch.cat([shuffle_rgb, shuffle_depth], 1).detach(), ori_iter) + patches = self.unfold2(shuffle_rgb).view(1, 3, self.dataset.kernel_size2 ** 2, shuffle_rgb.shape[-2], shuffle_rgb.shape[-1]) + else: + if (ori_iter > self.opt.ms_steps) and self.opt.use_another_mlp: + kernel_weights, mask = self.gaussians.mlp_rgb_ss(0, pos_enc, torch.cat([shuffle_rgb, shuffle_depth], 1).detach(), adj_iter) + patches = self.unfold_ss(shuffle_rgb).view(1, 3, self.dataset.kernel_size_ss ** 2, shuffle_rgb.shape[-2], shuffle_rgb.shape[-1]) + else: + kernel_weights, mask = self.gaussians.mlp_rgb_ms(0, pos_enc, torch.cat([shuffle_rgb, shuffle_depth], 1).detach(), ori_iter) + patches = self.unfold3(shuffle_rgb).view(1, 3, self.dataset.kernel_size3 ** 2, shuffle_rgb.shape[-2], shuffle_rgb.shape[-1]) + + kernel_weights = kernel_weights.unsqueeze(1) + rgb = torch.sum(patches * kernel_weights, 2)[0] + mask = mask[0] + blur_image = mask * rgb + (1 - mask) * image + + depthloss = self.opt.depth_loss_alpha * tv_loss(shuffle_depth) if self.opt.use_depth_loss else torch.tensor(0.0, device="cuda") + maskloss = self.opt.mask_loss_alpha * mask.mean() if self.opt.use_mask_loss else torch.tensor(0.0, device="cuda") + tvloss = self.opt.rgbtv_loss_alpha * tv_loss(shuffle_rgb) if self.opt.use_rgbtv_loss else torch.tensor(0.0, device="cuda") + + Ll1 = l1_loss(blur_image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(blur_image, gt_image)) + tvloss + maskloss + depthloss + else: + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + metrics = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + metrics[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if adj_iter < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if adj_iter > self.opt.densify_from_iter and adj_iter % self.opt.densification_interval == 0: + size_threshold = 20 if adj_iter > self.opt.opacity_reset_interval else None + if ori_iter <= self.opt.ms_steps: + dgt = self.opt.init_dgt + min_op = self.opt.init_opacity if self.opt.init_opacity >= 0 else self.opt.min_opacity + else: + dgt = self.opt.densify_grad_threshold + min_op = self.opt.min_opacity + self.gaussians.densify_and_prune(dgt, min_op, self.scene.cameras_extent, size_threshold) + + if ori_iter >= self.upsample_iter[0] and ori_iter < self.upsample_iter[1]: + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale2) + elif ori_iter >= self.upsample_iter[1]: + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale1) + else: + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale4) + + if adj_iter % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and adj_iter == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if adj_iter % 100 == 0 and adj_iter > self.opt.densify_until_iter: + if adj_iter < (self.opt.iterations - self.opt.ms_steps) - 100: + if ori_iter >= self.upsample_iter[0] and ori_iter < self.upsample_iter[1]: + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale2) + elif ori_iter >= self.upsample_iter[1]: + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale1) + else: + self.gaussians.compute_3D_filter(cameras=self.trainCameras_scale4) + + if ori_iter < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics.update({ + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + }) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_bgtriangle.py b/methods/wrapper_bgtriangle.py new file mode 100644 index 0000000000000000000000000000000000000000..d7b6a17f976544eec66a564c0e11a58c9ccc8e09 --- /dev/null +++ b/methods/wrapper_bgtriangle.py @@ -0,0 +1,261 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +def force_unload_conflicts(): + conflicts = ['render', 'scene', 'utils', 'arguments', 'model'] + for name in list(sys.modules.keys()): + parts = name.split('.') + if parts[0] in conflicts: + del sys.modules[name] + +@register_method("bgtriangle") +class BGTriangleWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + force_unload_conflicts() + + self.method_path = os.path.abspath("/root/autodl-tmp/bg_triangle_official") + if self.method_path not in sys.path: + sys.path.insert(0, self.method_path) + + try: + from utils.loss_utils import l1_loss, ssim + from arguments import ModelParams, PipelineParams, OptimizationParams + from model.bprimitive_bezier import BPrimitiveBezier + from render.renderer import Renderer + from scene import Scene + + self.l1_loss = l1_loss + self.ssim = ssim + + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.order = hyperparams.get("order", 1) + self.bprimitive = BPrimitiveBezier(self.order, self.dataset.sh_degree) + self.bprimitive.boundary_mode = 1 + + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.bprimitive) + self.bprimitive.training_setup(self.opt) + + N = self.bprimitive.control_point.shape[0] + if not hasattr(self.bprimitive, 'vis_map'): + self.bprimitive.vis_map = torch.zeros((N, 6, 6), device="cuda") + if not hasattr(self.bprimitive, 'vis_accum'): + self.bprimitive.vis_accum = torch.zeros((N, 1), device="cuda") + if not hasattr(self.bprimitive, 'denom_vis'): + self.bprimitive.denom_vis = torch.zeros((N, 1), device="cuda") + if not hasattr(self.bprimitive, 'edge_accum'): + self.bprimitive.edge_accum = torch.zeros((N, 1), device="cuda") + if not hasattr(self.bprimitive, 'denom_edge'): + self.bprimitive.denom_edge = torch.zeros((N, 1), device="cuda") + if not hasattr(self.bprimitive, 'gradient_accum'): + self.bprimitive.gradient_accum = torch.zeros((N, 1), device="cuda") + if not hasattr(self.bprimitive, 'denom'): + self.bprimitive.denom = torch.zeros((N, 1), device="cuda") + + self.renderer = Renderer() + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_primitives = self.bprimitive.control_point.shape[0] + + finally: + if self.method_path in sys.path: + sys.path.remove(self.method_path) + + def train_iteration(self, step): + force_unload_conflicts() + if self.method_path not in sys.path: + sys.path.insert(0, self.method_path) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg, debug_info = self.renderer( + self.bprimitive, viewpoint_cam, self.background, None, + self.pipe.num_segments_per_bprimitive_edge, -5.0, 1.0, "version_1" + ) + image = render_pkg + + gt_image = viewpoint_cam.original_image.cuda() + loss_l1 = self.l1_loss(image, gt_image) + loss_ssim = 1.0 - self.ssim(image, gt_image) + loss_target = (1.0 - 0.2) * loss_l1 + loss_parasitic = 0.2 * loss_ssim + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.bprimitive.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + # 使用 reshape 替换 view + grad_target = self.bprimitive.control_point.grad.clone().reshape(self.bprimitive.control_point.shape[0], -1) if self.bprimitive.control_point.grad is not None else torch.zeros_like(self.bprimitive.control_point).reshape(self.bprimitive.control_point.shape[0], -1) + + self.bprimitive.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + # 使用 reshape 替换 view + grad_parasitic = self.bprimitive.control_point.grad.clone().reshape(self.bprimitive.control_point.shape[0], -1) if self.bprimitive.control_point.grad is not None else torch.zeros_like(self.bprimitive.control_point).reshape(self.bprimitive.control_point.shape[0], -1) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.bprimitive.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.bprimitive.optimizer.step() + self.bprimitive.optimizer.zero_grad(set_to_none=True) + + num_primitives = self.bprimitive.control_point.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_primitives), + "delta_N": int(num_primitives - self.last_n_primitives), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_primitives = num_primitives + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.bprimitive.opacity).clone().detach() + scales = torch.exp(self.bprimitive.scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.bprimitive.control_point_dc.detach().norm(dim=-1).mean(dim=1) + + if self.method_path in sys.path: + sys.path.remove(self.method_path) + + return metrics, histograms + + def render(self, camera): + force_unload_conflicts() + if self.method_path not in sys.path: + sys.path.insert(0, self.method_path) + + with torch.no_grad(): + image, _ = self.renderer( + self.bprimitive, camera, self.background, None, + self.pipe.num_segments_per_bprimitive_edge, -5.0, 1.0, "version_1" + ) + + if self.method_path in sys.path: + sys.path.remove(self.method_path) + + return {"image": image, "depth": None} + + def save(self, save_dir, step): + force_unload_conflicts() + if self.method_path not in sys.path: + sys.path.insert(0, self.method_path) + + self.scene.save(step) + + if self.method_path in sys.path: + sys.path.remove(self.method_path) + + def load(self, model_path, iteration): + pass + + def get_spatial_centers(self): + return self.bprimitive.control_point.mean(dim=1) + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.bprimitive.scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.bprimitive.opacity))) + + dc = self.bprimitive.control_point_dc + rest = self.bprimitive.control_point_rest + if rest is not None and rest.shape[2] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.bprimitive.control_point.mean(dim=1) + opacities = torch.sigmoid(self.bprimitive.opacity).squeeze() + scales = torch.exp(self.bprimitive.scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_c3dgs.py b/methods/wrapper_c3dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..32820447014b46320246bf502dfa8f542ba0e332 --- /dev/null +++ b/methods/wrapper_c3dgs.py @@ -0,0 +1,299 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +_UPSTREAM = '/root/autodl-tmp/C3DGS_offy' +sys.path.insert(0, _UPSTREAM) +# monorepo-flat: insert each first-level subdir individually +import os as _os +for _sub in ('utils', 'scene', 'gaussian_renderer', 'arguments', 'lpipsPyTorch', 'experiments'): + _p = _os.path.join(_UPSTREAM, _sub) + if _os.path.isdir(_p): + sys.path.insert(0, _p) +del _os, _p, _sub, _UPSTREAM +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams, CompressionParams +from compression.vq import CompressionSettings, compress_gaussians + +@register_method("c3dgs") +class C3DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.cp = CompressionParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.comp = self.cp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree, quantization=not self.opt.not_quantization_aware) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians, load_iteration=-1, shuffle=True) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.compressed = False + self.step_offset = 30000 + + def calc_importance_internal(self): + scaling = self.gaussians.scaling_qa(self.gaussians.scaling_activation(self.gaussians._scaling.detach())) + cov3d = self.gaussians.covariance_activation(scaling, 1.0, self.gaussians.get_rotation.detach(), True).requires_grad_(True) + scaling_factor = self.gaussians.scaling_factor_activation(self.gaussians.scaling_factor_qa(self.gaussians._scaling_factor.detach())) + + h1 = self.gaussians._features_dc.register_hook(lambda grad: grad.abs()) + h2 = self.gaussians._features_rest.register_hook(lambda grad: grad.abs()) + h3 = cov3d.register_hook(lambda grad: grad.abs()) + + bg = torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device="cuda") + self.gaussians._features_dc.grad = None + self.gaussians._features_rest.grad = None + + num_pixels = 0 + for camera in self.scene.getTrainCameras(): + cov3d_scaled = cov3d * scaling_factor.square() + rendering = native_render(camera, self.gaussians, self.pipe, bg, clamp_color=False, cov3d=cov3d_scaled)["render"] + loss = rendering.sum() + loss.backward() + num_pixels += rendering.shape[1] * rendering.shape[2] + + importance = torch.cat([self.gaussians._features_dc.grad, self.gaussians._features_rest.grad], 1).flatten(-2) / num_pixels + cov_grad = cov3d.grad / num_pixels + + h1.remove() + h2.remove() + h3.remove() + return importance.detach(), cov_grad.detach() + + def train_iteration(self, step): + if not self.compressed: + color_importance, gaussian_sensitivity = self.calc_importance_internal() + color_importance_n = color_importance.amax(-1) + gaussian_importance_n = gaussian_sensitivity.amax(-1) + + c_set = CompressionSettings( + codebook_size=self.comp.color_codebook_size, + importance_prune=self.comp.color_importance_prune, + importance_include=self.comp.color_importance_include, + steps=int(self.comp.color_cluster_iterations), + decay=self.comp.color_decay, + batch_size=self.comp.color_batch_size, + ) + g_set = CompressionSettings( + codebook_size=self.comp.gaussian_codebook_size, + importance_prune=None, + importance_include=self.comp.gaussian_importance_include, + steps=int(self.comp.gaussian_cluster_iterations), + decay=self.comp.gaussian_decay, + batch_size=self.comp.gaussian_batch_size, + ) + compress_gaussians(self.gaussians, color_importance_n, gaussian_importance_n, c_set, g_set, self.comp.color_compress_non_dir, prune_threshold=self.comp.prune_threshold) + self.compressed = True + self.gaussians.training_setup(self.opt) + + actual_step = step + self.step_offset + self.gaussians.update_learning_rate(actual_step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation, self.gaussians._scaling_factor], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest] + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + os.makedirs(os.path.join(save_dir, "point_cloud", f"iteration_{step}"), exist_ok=True) + self.gaussians.save_npz(os.path.join(save_dir, "point_cloud", f"iteration_{step}", "point_cloud.npz")) + + def load(self, model_path, iteration): + self.gaussians.load_npz(os.path.join(model_path, "point_cloud", f"iteration_{iteration}", "point_cloud.npz")) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_cdcgs.py b/methods/wrapper_cdcgs.py new file mode 100644 index 0000000000000000000000000000000000000000..f7da3d9362f8b7e4b46c2f279a19a3d40357d405 --- /dev/null +++ b/methods/wrapper_cdcgs.py @@ -0,0 +1,278 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../CDC-GS_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from gaussian_renderer import render_cdc +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def update_freq_cuda_local(views, gaussians, pipe, bg): + max_freqs = torch.zeros_like(gaussians.get_opacity, dtype=torch.float, device="cuda") + for view in views: + if not hasattr(view, 'mask_H'): + view.mask_H = torch.ones((1, 1, view.image_height, view.image_width), device="cuda") + project_pkg = render_cdc(view, gaussians, pipe, bg) + freq = project_pkg["test_freq"].unsqueeze(1) + torch.maximum(max_freqs, freq, out=max_freqs) + gaussians.max_freqs = max_freqs + +@register_method("cdcgs") +class CDCGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.gs_num = 9999999 + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + self.gaussians.max_freqs = torch.zeros((self.gaussians.get_xyz.shape[0], 1), dtype=torch.float, device="cuda") + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + if not hasattr(viewpoint_cam, 'mask_H'): + viewpoint_cam.mask_H = torch.ones((1, 1, viewpoint_cam.image_height, viewpoint_cam.image_width), device="cuda") + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gt_image = viewpoint_cam.original_image.cuda() + + loss_target = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_parasitic = 1.0 - ssim_value + loss = (1.0 - self.opt.lambda_dssim) * loss_target + self.opt.lambda_dssim * loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + lr = 0.001 + for pg in self.gaussians.optimizer.param_groups: + if pg["params"][0] is p: + lr = pg["lr"] + break + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + group_grads.append(u.reshape(-1)) + if group_grads: + grads_target[group_name] = torch.cat(group_grads) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + lr = 0.001 + for pg in self.gaussians.optimizer.param_groups: + if pg["params"][0] is p: + lr = pg["lr"] + break + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + group_grads.append(u.reshape(-1)) + if group_grads: + grads_parasitic[group_name] = torch.cat(group_grads) + + for group_name in param_groups_map: + gt = grads_target.get(group_name) + gp = grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + update_freq_cuda_local(self.scene.getTrainCameras().copy(), self.gaussians, self.pipe, self.background) + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii, max_num=self.args.gs_num, iteration=step, views=self.scene.getTestCameras().copy(), pipe=self.pipe, bg=self.background) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "freq_mean": float(self.gaussians.max_freqs.mean()) if self.gaussians.max_freqs.numel() > 0 else 0.0 + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + if not hasattr(camera, 'mask_H'): + camera.mask_H = torch.ones((1, 1, camera.image_height, camera.image_width), device="cuda") + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_coadaptgs.py b/methods/wrapper_coadaptgs.py new file mode 100644 index 0000000000000000000000000000000000000000..dde5b8db181ba948d913b7fd3d7bcde35242607b --- /dev/null +++ b/methods/wrapper_coadaptgs.py @@ -0,0 +1,305 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../CoAdaptGS'))) +from utils.loss_utils import l1_loss, ssim, SmoothLoss +from utils.graphics_utils import inverse_warp_images +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("coadaptgs") +class CoAdaptGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + + self.parser.add_argument("--opacity_decay", action="store_true", default=True) + self.parser.add_argument("--opacity_decay_factor", type=float, default=0.995) + self.parser.add_argument("--dropout_factor", type=float, default=0.2) + self.parser.add_argument("--sigma_noise", type=float, default=0.8) + self.parser.add_argument("--cam_trans_dist", type=float, default=0.4) + self.parser.add_argument("--binocular_consistency", action="store_true", default=True) + self.parser.add_argument("--shift_cam_start", type=int, default=20000) + self.parser.add_argument("--n_views", type=int, default=-1) + self.parser.add_argument("--dataset_name", type=str, default="") + self.parser.add_argument("--suffix", type=str, default="") + + self.args, _ = self.parser.parse_known_args() + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + # SplatAtlas: required by CoAdaptGS scene/__init__.py but not in ModelParams defaults + self.args.n_views = 0 + self.args.dataset_name = "LLFF" + self.args.suffix = None + if not hasattr(self.args, 'images') or self.args.images is None: + self.args.images = "images" + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.smooth_loss = SmoothLoss() + + def train_iteration(self, step): + N = self.gaussians.get_xyz.shape[0] + if self.gaussians._scaling.shape[0] != N: + fallback_scale = float(self.gaussians._scaling.mean().item()) if self.gaussians._scaling.numel() > 0 else -4.0 + new_scales = torch.ones((N, 3), dtype=torch.float32, device="cuda") * fallback_scale + self.gaussians._scaling = torch.nn.Parameter(new_scales.requires_grad_(True)) + self.gaussians.training_setup(self.opt) + + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = self.background + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, + dropout_factor=self.args.dropout_factor, + sigma_noise=self.args.sigma_noise, train=True) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + depth = render_pkg.get("rendered_depth", render_pkg.get("depth", torch.zeros_like(image[0:1]))) + alpha = render_pkg.get("rendered_alpha", torch.zeros_like(image[0:1])) + dropout_mask = render_pkg.get("dropout_mask", torch.ones(self.gaussians.get_xyz.shape[0], dtype=torch.bool, device="cuda")) + + gt_image = viewpoint_cam.original_image.cuda() + + disparity_loss = torch.tensor(0.0, device="cuda") + if self.args.binocular_consistency and step > self.args.shift_cam_start: + trans_dist = (torch.rand(1) * self.args.cam_trans_dist * random.choice([-1.0, 1.0])).item() + shifted_cam = self.scene.getShiftedCamera(viewpoint_cam, trans_dist) + shifted_pkg = native_render(shifted_cam, self.gaussians, self.pipe, bg, dropout_factor=0.0, sigma_noise=0.0, train=True) + shifted_image = shifted_pkg["render"] + + focal_x = viewpoint_cam.get_focal()[0] if hasattr(viewpoint_cam, 'get_focal') else viewpoint_cam.focal_x + disparity = focal_x * (-trans_dist) / (depth + 1e-5) + + image_height = viewpoint_cam.image_height + image_width = viewpoint_cam.image_width + row_indices = torch.arange(0, image_height).view(-1, 1).repeat(1, image_width).cuda() + column_indices = torch.arange(0, image_width).repeat(image_height, 1).cuda() + mask_tensor = torch.ones((1, image_height, image_width), dtype=torch.float32).cuda() + + warped_image = inverse_warp_images(shifted_image.unsqueeze(0), disparity.unsqueeze(0), row_indices, column_indices) + shift_mask = inverse_warp_images(mask_tensor.unsqueeze(0), disparity.unsqueeze(0), row_indices, column_indices) + disparity_loss = l1_loss(warped_image, gt_image.unsqueeze(0), mask=shift_mask) + 0.05 * self.smooth_loss(disparity=disparity*shift_mask, image=gt_image.unsqueeze(0)) + + alpha_loss = torch.tensor(0.0, device="cuda") + if hasattr(viewpoint_cam, "gt_alpha_mask") and viewpoint_cam.gt_alpha_mask is not None: + alpha_loss = torch.mean(torch.abs(alpha) * (1 - viewpoint_cam.gt_alpha_mask)) + + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + disparity_loss + alpha_loss + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + eta_t = self.gaussians.optimizer.param_groups[0]["lr"] + u_target = (eta_t / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_parasitic = (eta_t / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + else: + u_target = grad_target + u_parasitic = grad_parasitic + + valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if self.args.opacity_decay: + self.gaussians.opacity_decay(factor=self.args.opacity_decay_factor) + + if step < self.opt.densify_until_iter: + true_indices = torch.nonzero(dropout_mask, as_tuple=True)[0] + filtered_indices = true_indices[visibility_filter] + combined_mask = torch.zeros_like(dropout_mask, dtype=torch.bool).cuda() + combined_mask[filtered_indices] = True + self.gaussians.max_radii2D[combined_mask] = torch.max(self.gaussians.max_radii2D[combined_mask], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, combined_mask) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + dropout_survival_ratio = float(dropout_mask.float().mean()) + visible_area_ratio = float(render_pkg.get("visi_area", torch.zeros(1)).float().mean()) + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic - disparity_loss - alpha_loss), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "dropout_survival_ratio": dropout_survival_ratio, + "disparity_loss": float(disparity_loss), + "visible_area_ratio": visible_area_ratio + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, dropout_factor=0.0, sigma_noise=0.0, train=False) + + image = render_pkg["render"] + depth = render_pkg.get("rendered_depth", render_pkg.get("depth", None)) + + if depth is None: + depth = torch.zeros_like(image[0:1]) + + if depth.dim() == 3: + depth = depth.unsqueeze(0) + + return {"image": image, "depth": depth} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_compact3dgs.py b/methods/wrapper_compact3dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..088dc553c1b98dd5824725c3c6f344a8b6e3a581 --- /dev/null +++ b/methods/wrapper_compact3dgs.py @@ -0,0 +1,280 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +_UPSTREAM = '/root/autodl-tmp/Compact3DGS_offy' +sys.path.insert(0, _UPSTREAM) +# monorepo-flat: insert each first-level subdir individually +import os as _os +for _sub in ('utils', 'scene', 'gaussian_renderer', 'arguments', 'lpipsPyTorch', 'experiments'): + _p = _os.path.join(_UPSTREAM, _sub) + if _os.path.isdir(_p): + sys.path.insert(0, _p) +del _os, _p, _sub, _UPSTREAM +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("compact3dgs") +class Compact3DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + rvq_iter_flag = step > self.opt.rvq_iter + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, itr=step, rvq_iter=rvq_iter_flag) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + loss_parasitic = self.opt.lambda_mask * torch.mean(torch.sigmoid(self.gaussians._mask)) + loss = loss_target + loss_parasitic + + stats = {} + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity, self.gaussians._mask], + "appearance": [self.gaussians.recolor.params, self.gaussians.mlp_head.params] + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_net.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for g_name, params in param_groups_map.items(): + opt = self.gaussians.optimizer_net if g_name == "appearance" else self.gaussians.optimizer + valid_grads = [] + for p in params: + if p.grad is not None: + state = opt.state.get(p, {}) + v_t = state.get("exp_avg_sq", torch.ones_like(p)) + lr = 0.001 + for pg in opt.param_groups: + if id(p) in [id(param) for param in pg["params"]]: + lr = pg["lr"] + break + u = (lr / (torch.sqrt(v_t) + 1e-8)) * p.grad.clone() + valid_grads.append(u.reshape(-1)) + grads_target[g_name] = torch.cat(valid_grads) if valid_grads else torch.empty(0, device="cuda") + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_net.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for g_name, params in param_groups_map.items(): + opt = self.gaussians.optimizer_net if g_name == "appearance" else self.gaussians.optimizer + valid_grads = [] + for p in params: + if p.grad is not None: + state = opt.state.get(p, {}) + v_t = state.get("exp_avg_sq", torch.ones_like(p)) + lr = 0.001 + for pg in opt.param_groups: + if id(p) in [id(param) for param in pg["params"]]: + lr = pg["lr"] + break + u = (lr / (torch.sqrt(v_t) + 1e-8)) * p.grad.clone() + valid_grads.append(u.reshape(-1)) + grads_parasitic[g_name] = torch.cat(valid_grads) if valid_grads else torch.empty(0, device="cuda") + + for g_name in param_groups_map: + gt = grads_target.get(g_name) + gp = grads_parasitic.get(g_name) + if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{g_name}"] = ti + + flat_gt = torch.cat([g for g in grads_target.values() if g.numel() > 0]) + flat_gp = torch.cat([g for g in grads_parasitic.values() if g.numel() > 0]) + if flat_gt.numel() > 0 and flat_gp.numel() > 0 and flat_gt.norm() > 0 and flat_gp.norm() > 0: + grad_cos_sim = float(F.cosine_similarity(flat_gt.unsqueeze(0), flat_gp.unsqueeze(0))) + parasitic_ratio = float(flat_gp.norm() / (flat_gt.norm() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_net.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + else: + if step % self.opt.mask_prune_iter == 0: + self.gaussians.mask_prune() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_net.step() + self.gaussians.optimizer_net.zero_grad(set_to_none=True) + self.gaussians.scheduler_net.step() + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = self.gaussians.get_scaling.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["mask_activation"] = torch.sigmoid(self.gaussians._mask).clone().detach() + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, itr=-1, rvq_iter=True) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + comp = True + store_npz = True + if step == self.opt.iterations: + self.gaussians.final_prune(compress=comp) + self.gaussians.precompute() + self.scene.save(step, compress=comp, store=store_npz) + + def load(self, model_path, iteration): + load_path = os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud') + self.gaussians.load_model(load_path) + self.gaussians.precompute() + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + metrics["mask_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._mask))) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians.get_rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() * torch.sigmoid(self.gaussians._mask).squeeze() + scales = self.gaussians.get_scaling + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_compgs.py b/methods/wrapper_compgs.py new file mode 100644 index 0000000000000000000000000000000000000000..6b6146025d27d793b62a577947588063ac4ea6e4 --- /dev/null +++ b/methods/wrapper_compgs.py @@ -0,0 +1,214 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../CompGS_offy'))) +from Modules.TrainerCompGS import TrainerCompGS +from Modules.Common import RenderSettings +from pytorch_msssim import ssim + +@register_method("compgs") +class CompGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.source_path = dataset_config["source_path"] + self.model_path = dataset_config["model_path"] + self.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + dummy_config = os.path.join(os.path.dirname(__file__), '../../CompGS_offy/Configs/TanksAndTemplates.yaml') + import yaml + base_cfg_path = "/root/autodl-tmp/CompGS_offy/Configs/TanksAndTemplates.yaml" + with open(base_cfg_path, "r") as f: + cfg = yaml.safe_load(f) + + # 强制接管核心路径与参数 + cfg["dataset"]["root"] = dataset_config["source_path"] + cfg["training"]["save_directory"] = dataset_config["model_path"] + + # 动态对齐迭代步数 + iters = hyperparams.get("iterations", 30000) + cfg["training"]["max_iterations"] = iters + + os.makedirs(dataset_config["model_path"], exist_ok=True) + runtime_cfg_path = os.path.join(dataset_config["model_path"], "runtime_config.yaml") + with open(runtime_cfg_path, "w") as f: + yaml.dump(cfg, f) + + self.trainer = TrainerCompGS(config_path=runtime_cfg_path, override_cfgs={}) + self.gaussian_model = self.trainer.gaussian_model + self.dataset = self.trainer.dataset + self.optimizer = self.trainer.gaussian_optimizer + self.aux_optimizer = self.trainer.aux_optimizer + + self.last_n_gaussians = self.gaussian_model.num_coupled_primitive + + def train_iteration(self, step): + self.trainer.gaussian_lr_scheduler(iteration=step) + + sample = self.dataset[random.randint(0, len(self.dataset) - 1)] + + render_settings = RenderSettings( + cam_idx=sample.cam_idx, image_height=sample.image_height, image_width=sample.image_width, + tanfovx=sample.tan_half_fov_x, tanfovy=sample.tan_half_fov_y, campos=sample.camera_center, + viewmatrix=sample.world_to_view_proj_mat, projmatrix=sample.world_to_image_proj_mat) + + retain_grad = step < self.trainer.configs['adaptive_control']['stop_iteration'] + render_results = self.gaussian_model.render(render_settings=render_settings, retain_grad=retain_grad) + + ssim_weight = self.trainer.configs['training']['ssim_weight'] + l1_loss = F.l1_loss(render_results.rendered_img, sample.img) + ssim_loss = 1 - ssim(render_results.rendered_img.unsqueeze(dim=0), sample.img.unsqueeze(dim=0), data_range=1., size_average=True) + rendering_loss = (1 - ssim_weight) * l1_loss + ssim_weight * ssim_loss + + reg_loss = 0.01 * render_results.scales.prod(dim=1).mean() + bpp = render_results.bpp + rate_loss = self.trainer.configs['training']['lambda_weight'] * sum(v for v in bpp.values()) if step > self.trainer.configs['training']['rate_loss_start_iteration'] else torch.tensor(0.0, device=self.trainer.device) + + loss_target = rendering_loss + loss_parasitic = reg_loss + rate_loss + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + params = self.gaussian_model.gaussian_params + param_groups_map = { + "spatial": [params.means], + "geometry": [params.scales_before_exp, params.rotations_before_norm], + "opacity": [params.ref_feats], + "appearance": [params.res_feats], + } + + self.optimizer.zero_grad(set_to_none=True) + self.aux_optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + grad_target = params.means.grad.clone() if params.means.grad is not None else torch.zeros_like(params.means) + grads_target = {} + for group_name, plist in param_groups_map.items(): + grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in plist if p.grad is not None]) + + self.optimizer.zero_grad(set_to_none=True) + self.aux_optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + grad_parasitic = params.means.grad.clone() if params.means.grad is not None else torch.zeros_like(params.means) + grads_parasitic = {} + for group_name, plist in param_groups_map.items(): + grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in plist if p.grad is not None]) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.optimizer.zero_grad(set_to_none=True) + self.aux_optimizer.zero_grad(set_to_none=True) + loss.backward() + aux_loss = self.gaussian_model.aux_loss + aux_loss.backward() + else: + loss.backward() + aux_loss = self.gaussian_model.aux_loss + aux_loss.backward() + + self.trainer.optimize(iteration=step, render_results=render_results) + + num_gaussians = self.gaussian_model.num_coupled_primitive + metrics = { + "loss": float(loss), "loss_l1": float(l1_loss), "loss_ssim": float(ssim_loss), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["scaling"] = torch.exp(self.gaussian_model.gaussian_params.scales_before_exp).clone().detach() + scales_2d = histograms["scaling"][:, :2] + histograms["anisotropy"] = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["sh_dc_mag"] = self.gaussian_model.gaussian_params.ref_feats.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_settings = RenderSettings( + cam_idx=camera.cam_idx, image_height=camera.image_height, image_width=camera.image_width, + tanfovx=camera.tan_half_fov_x, tanfovy=camera.tan_half_fov_y, campos=camera.camera_center, + viewmatrix=camera.world_to_view_proj_mat, projmatrix=camera.world_to_image_proj_mat) + rendered_img, _, _ = self.gaussian_model.render_inference(render_settings=render_settings) + return {"image": rendered_img, "depth": None} + + def save(self, save_dir, step): + ckpt_folder = os.path.join(save_dir, f'iteration_{step}') + os.makedirs(ckpt_folder, exist_ok=True) + self.gaussian_model.save_uncompressed_params(os.path.join(ckpt_folder, 'point_cloud.ply')) + self.gaussian_model.save_weights(os.path.join(ckpt_folder, 'weights.pth')) + + def load(self, model_path, iteration): + ckpt_folder = os.path.join(model_path, f'iteration_{iteration}') + self.gaussian_model.load_uncompressed_params(os.path.join(ckpt_folder, 'point_cloud.ply')) + self.gaussian_model.load_weights(os.path.join(ckpt_folder, 'weights.pth')) + + def get_spatial_centers(self): + return self.gaussian_model.means + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.gaussian_model.gaussian_params.scales_before_exp) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + + ref_f = self.gaussian_model.gaussian_params.ref_feats + res_f = self.gaussian_model.gaussian_params.res_feats + if res_f is not None and res_f.shape[1] > 0: + metrics["sh_energy_ratio"] = float(res_f.norm(dim=-1).mean() / (ref_f.norm(dim=-1).mean() + 1e-7)) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device=self.trainer.device) + xyz = self.gaussian_model.means + scales = torch.exp(self.gaussian_model.gaussian_params.scales_before_exp) + sigma_sq = scales[:, :2].max(dim=1)[0].pow(2) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30000000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights, dim=1) + return densities + diff --git a/methods/wrapper_conegs.py b/methods/wrapper_conegs.py new file mode 100644 index 0000000000000000000000000000000000000000..d6f6d081448130db53b9c57baed5f73496daa29e --- /dev/null +++ b/methods/wrapper_conegs.py @@ -0,0 +1,293 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from omegaconf import OmegaConf +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../conegs'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.nerf_model import NeRFModel +from utils.nerf_utils import get_num_rays_to_cast + +@register_method("conegs") +class ConeGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + conegs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../conegs')) + if not GlobalHydra().is_initialized(): + initialize_config_dir(config_dir=os.path.join(conegs_dir, 'configs'), version_base=None) + self.cfg = compose(config_name="defaults") + + OmegaConf.set_struct(self.cfg, False) + + self.cfg.gaussian_model.source_path = dataset_config["source_path"] + self.cfg.gaussian_model.model_path = dataset_config["model_path"] + OmegaConf.set_struct(self.cfg, False) + self.cfg.gaussian_model.eval = True + if hasattr(self.cfg, 'dataset'): + self.cfg.dataset.eval = True + self.cfg.optimization.resolution = float(dataset_config.get("resolution", 1)) + + self.cfg.nerf_model.num_iters_pretrain = 20000 + self.cfg.nerf_model.nerf_train_during_3dgs = False + self.cfg.optimization.use_preactivation_opacity_penalty = True + self.cfg.optimization.opacity_penalty = 0.0002 + + if "max_points" in hyperparams: + self.cfg.optimization.max_points = hyperparams["max_points"] + + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.gaussians = GaussianModel(self.cfg.gaussian_model.sh_degree) + self.scene = Scene(self.cfg.gaussian_model, self.gaussians, stack_train_images=self.cfg.nerf_model.stack_images) + self.gaussians.training_setup(self.cfg.optimization) + + bg_color = [1, 1, 1] if self.cfg.gaussian_model.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + + self.nerf_model = NeRFModel(self.cfg.nerf_model) + self.nerf_model.pretrain_model(self.cfg, self.gaussians, self.scene, self.background) + + self.nerf_model.prop_optimizer.zero_grad(set_to_none=True) + self.nerf_model.model_optimizer.zero_grad(set_to_none=True) + if hasattr(self.scene, "stacked_images"): + del self.scene.stacked_images + if hasattr(self.scene, "stacked_cam2worlds"): + del self.scene.stacked_cam2worlds + if hasattr(self.scene, "stacked_pix2cams"): + del self.scene.stacked_pix2cams + + for i in range(len(self.scene.getTrainCameras())): + self.scene.getTrainCameras()[i].original_image = self.scene.getTrainCameras()[i].original_image.cuda() + for i in range(len(self.scene.getTestCameras())): + self.scene.getTestCameras()[i].original_image = self.scene.getTestCameras()[i].original_image.cuda() + + num_rays_to_cast = get_num_rays_to_cast(len(self.gaussians._xyz), gaussian_percentage_increase=self.cfg.optimization.gaussian_percentage_increase, nerf_train_during_3dgs=self.cfg.nerf_model.nerf_train_during_3dgs) + self.nerf_model.estimator_params["n_rays"] = num_rays_to_cast + self.nerf_model.estimator_params["stratified"] = False + self.nerf_model.estimator_params["requires_grad"] = False + + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_u_step(self, params, grads): + u_vecs = [] + for p, g in zip(params, grads): + if g is None: + continue + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + pg_lr = 1e-4 + for pg in self.gaussians.optimizer.param_groups: + if any(p is param for param in pg["params"]): + pg_lr = pg["lr"] + break + u = (pg_lr / (torch.sqrt(v) + 1e-15)) * g + else: + u = g + u_vecs.append(u.view(-1)) + if len(u_vecs) == 0: + return torch.zeros(1, device="cuda") + return torch.cat(u_vecs) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % self.cfg.optimization.SH_increase_iter == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + gt_image = viewpoint_cam.original_image.cuda() + + bg = torch.rand((3), device="cuda") if self.cfg.optimization.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.cfg.pipeline, bg) + image = render_pkg["render"] + + l1_full = torch.abs((image - gt_image)) + ssim_full = ssim(image, gt_image) + loss_target = (1.0 - self.cfg.optimization.lambda_dssim) * l1_full.mean() + loss_ssim_part = self.cfg.optimization.lambda_dssim * (1.0 - ssim_full.mean()) + + opacity = self.gaussians._opacity if self.cfg.optimization.use_preactivation_opacity_penalty else torch.abs(self.gaussians.get_opacity) + loss_opacity = self.cfg.optimization.opacity_penalty * opacity.mean() + + loss = loss_target + loss_ssim_part + loss_opacity + + weights = (1.0 - self.cfg.optimization.densification_lambda_dssim) * l1_full.mean(0) + self.cfg.optimization.densification_lambda_dssim * (1.0 - ssim_full.mean(0)) + weights = torch.clamp_min(weights, 0.0) + + metrics = {} + + if self.track_decoupling and step % 100 == 0: + semantic_groups = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest] + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target_dict = {} + for k, params in semantic_groups.items(): + grad_target_dict[k] = [p.grad.clone() if p.grad is not None else torch.zeros_like(p) for p in params] + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_opacity.backward(retain_graph=True) + grad_parasitic_dict = {} + for k, params in semantic_groups.items(): + grad_parasitic_dict[k] = [p.grad.clone() if p.grad is not None else torch.zeros_like(p) for p in params] + + for k in semantic_groups.keys(): + u_t = self._get_u_step(semantic_groups[k], grad_target_dict[k]) + u_p = self._get_u_step(semantic_groups[k], grad_parasitic_dict[k]) + + norm_t = torch.norm(u_t) + norm_p = torch.norm(u_p) + + if norm_p > 1e-7 and norm_t > 1e-7: + s_group = F.cosine_similarity(u_t.unsqueeze(0), u_p.unsqueeze(0)).squeeze() + r_group = norm_p / (norm_t + norm_p + 1e-7) + ti = float(r_group * torch.clamp(-s_group, min=0.0)) + else: + ti = 0.0 + metrics[f"TI_{k}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if self.cfg.optimization.use_3dgs_densification and step < self.cfg.optimization.densify_until_iter: + radii = render_pkg["radii"] + visibility_filter = render_pkg["visibility_filter"] + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg.get("optimized_samples_index", 0)) + + if step > self.cfg.optimization.densify_from_iter and step % 100 == 0: + size_threshold = 20 if step > 3000 else None + self.gaussians.densify_and_prune(0.0002, 0.005, self.scene.cameras_extent, size_threshold, max_points=self.cfg.optimization.max_points) + + if step % 3000 == 0: + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + if self.cfg.optimization.densify_from_iter < step and step < self.cfg.optimization.densify_until_iter and (self.cfg.optimization.gaussian_percentage_increase or self.cfg.optimization.max_points): + num_rays_to_cast = get_num_rays_to_cast(len(self.gaussians._xyz), gaussian_percentage_increase=self.cfg.optimization.gaussian_percentage_increase, nerf_train_during_3dgs=self.cfg.nerf_model.nerf_train_during_3dgs) + self.nerf_model.cast_rays_during_optimization(weights, num_rays_to_cast, viewpoint_cam, self.cfg, self.gaussians, self.scene, bg, step) + + with torch.no_grad(): + if self.cfg.optimization.densify_from_iter < step and step % 100 == 0: + dead_mask = (self.gaussians.get_opacity <= 0.005).squeeze(-1) + if step < self.cfg.optimization.densify_until_iter: + self.gaussians.prune_points(dead_mask) + if self.cfg.optimization.gaussian_percentage_increase or self.cfg.optimization.max_points: + self.gaussians.initialize_new_points(use_cone_radius=True, max_points=self.cfg.optimization.max_points, opacity_init_value=self.cfg.optimization.opacity_init_value) + + num_gaussians = self.gaussians.get_xyz.shape[0] + stashed_count = len(self.gaussians.stashed_xyz) if hasattr(self.gaussians, "stashed_xyz") and self.gaussians.stashed_xyz is not None else 0 + + metrics.update({ + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_ssim_part), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "stashed_points_count": int(stashed_count), + "penalty_loss_ratio": float(loss_opacity / (loss + 1e-7)) + }) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.cfg.pipeline, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_contextgs.py b/methods/wrapper_contextgs.py new file mode 100644 index 0000000000000000000000000000000000000000..90a81a59b09b7240e8a9d1e74c1c9cfd4e608677 --- /dev/null +++ b/methods/wrapper_contextgs.py @@ -0,0 +1,316 @@ +import os +import sys +import random +import math +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../ContextGS_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import prefilter_voxel, render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("contextgs") +class ContextGS_offyWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + + self.parser.add_argument('--level_num', type=int, default=3) + self.parser.add_argument('--level_scale', type=int, default=10) + self.parser.add_argument("--n_features", type=int, default=4) + self.parser.add_argument("--lmbda", type=float, default=0.001) + self.parser.add_argument("--lmbda_rec", type=float, default=1) + self.parser.add_argument("--disable_hyper", default=False, action="store_true") + + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.args_param = self.args + + self.gaussians = GaussianModel( + self.dataset.feat_dim, + self.dataset.n_offsets, + self.dataset.voxel_size, + self.dataset.update_depth, + self.dataset.update_init_factor, + self.dataset.update_hierachy_factor, + self.dataset.use_feat_bank, + n_features_per_level=self.args_param.n_features, + level_num=self.args_param.level_num, + hyper_divisor=self.dataset.hyper_divisor, + target_ratio=self.dataset.target_ratio, + disable_hyper=self.args_param.disable_hyper + ) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.update_anchor_bound() + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = self.gaussians.get_anchor.shape[0] + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + retain_grad = (step < self.opt.update_until and step >= 0) + voxel_visible_mask = prefilter_voxel(viewpoint_cam, self.gaussians, self.pipe, self.background) + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask, retain_grad=retain_grad, step=step) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + offset_selection_mask = render_pkg["selection_mask"] + opacity = render_pkg["neural_opacity"] + scaling = render_pkg["scaling"] + bit_per_param = render_pkg.get("bit_per_param", None) + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + scaling_reg = scaling.prod(dim=1).mean() + + loss_target = self.args_param.lmbda_rec * ((1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value)) + + loss_parasitic = 0.01 * scaling_reg + if bit_per_param is not None: + loss_parasitic = loss_parasitic + self.args_param.lmbda * bit_per_param + loss_parasitic = loss_parasitic + 5e-4 * torch.mean(torch.sigmoid(self.gaussians._mask)) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + param_groups_map = { + "spatial": [self.gaussians._anchor, self.gaussians._offset], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity, self.gaussians._mask], + "appearance": [self.gaussians._anchor_feat, self.gaussians._hyper_latent], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = [] + for p in params: + if p.grad is not None: + grads_target[group_name].append(p.grad.clone()) + else: + grads_target[group_name].append(torch.zeros_like(p)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = [] + for p in params: + if p.grad is not None: + grads_parasitic[group_name].append(p.grad.clone()) + else: + grads_parasitic[group_name].append(torch.zeros_like(p)) + + for group_name, params in param_groups_map.items(): + u_t_list = [] + u_p_list = [] + for i, p in enumerate(params): + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + else: + v_t = torch.ones_like(p) + + for param_group in self.gaussians.optimizer.param_groups: + if id(p) in [id(opt_p) for opt_p in param_group['params']]: + lr = param_group['lr'] + break + else: + lr = 1e-3 + + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grads_target[group_name][i] + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grads_parasitic[group_name][i] + u_t_list.append(u_t.reshape(-1)) + u_p_list.append(u_p.reshape(-1)) + + gt = torch.cat(u_t_list) + gp = torch.cat(u_p_list) + + if gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + all_gt = torch.cat([torch.cat(u_t_list) for group_name, params in param_groups_map.items() for i, p in enumerate(params)]) + all_gp = torch.cat([torch.cat(u_p_list) for group_name, params in param_groups_map.items() for i, p in enumerate(params)]) + if all_gt.norm() > 0 and all_gp.norm() > 0: + grad_cos_sim = float(F.cosine_similarity(all_gt.unsqueeze(0), all_gp.unsqueeze(0))) + parasitic_ratio = float(all_gp.norm() / (all_gt.norm() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.update_until and step > self.opt.start_stat: + self.gaussians.training_statis(viewspace_point_tensor, opacity, visibility_filter, offset_selection_mask, voxel_visible_mask) + if step not in range(3000, 4000): + if step > self.opt.update_from and step % self.opt.update_interval == 0: + self.gaussians.adjust_anchor(check_interval=self.opt.update_interval, success_threshold=self.opt.success_threshold, grad_threshold=self.opt.densify_grad_threshold, min_opacity=self.opt.min_opacity) + elif step == self.opt.update_until: + del self.gaussians.opacity_accum + del self.gaussians.offset_gradient_accum + del self.gaussians.offset_denom + torch.cuda.empty_cache() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_anchor.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + + if bit_per_param is not None: + metrics["bit_per_param"] = float(bit_per_param) + + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = self.gaussians.get_scaling.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["anchor_feat_mag"] = self.gaussians._anchor_feat.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + voxel_visible_mask = prefilter_voxel(camera, self.gaussians, self.pipe, self.background) + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply_sparse_gaussian(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians.get_anchor + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + feat, hyper = self.gaussians._anchor_feat, self.gaussians._hyper_latent + if hyper is not None and hyper.shape[1] > 0: + metrics["hyper_energy_ratio"] = float(hyper.norm(dim=-1).mean() / (feat.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians.get_anchor + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = self.gaussians.get_scaling + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + diff --git a/methods/wrapper_drkgs.py b/methods/wrapper_drkgs.py new file mode 100644 index 0000000000000000000000000000000000000000..3d671da7a5486da17391d5d8a525dce8e6571b13 --- /dev/null +++ b/methods/wrapper_drkgs.py @@ -0,0 +1,254 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../DRK_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import drk_render_func +from scene import Scene +from scene.gaussian_model import DRKModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("drkgs") +class DRKGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.gs_type = 'DRK' + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + if "360" in str(self.args.source_path): + self.opt.is_unbounded = True + self.opt.kernel_density = "dense" + + self.gaussians = DRKModel(self.dataset.sh_degree, training=True) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_effective_step(self, params): + vec = [] + for p in params: + if p.grad is not None: + g = p.grad.clone() + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get("exp_avg_sq", torch.ones_like(p) * 1e-8) + lr = 1e-4 + for pg in self.gaussians.optimizer.param_groups: + if len(pg["params"]) > 0 and pg["params"][0] is p: + lr = pg["lr"] + break + u = (lr / (torch.sqrt(v) + 1e-8)) * g + vec.append(u.reshape(-1)) + return torch.cat(vec) if vec else torch.zeros(1, device="cuda") + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + self.gaussians.update(step) + + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = drk_render_func(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points_densify"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + opacity_grad_densify = render_pkg.get("opacity_grad_densify", None) + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation, self.gaussians._thetas], + "opacity": [self.gaussians._opacity, self.gaussians._acutance], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest, self.gaussians._l1l2_rates], + } + + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = self._get_effective_step(params) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = self._get_effective_step(params) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter, opacity_tensor=opacity_grad_densify) + + if step > self.opt.densify_from_iter and step % self.gaussians.densification_interval == 0: + size_threshold = 20 if step > self.gaussians.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.gaussians.densify_grad_threshold, self.gaussians.min_opacity_pruning, self.scene.cameras_extent, size_threshold) + + if step != 0 and step % self.gaussians.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = self.gaussians.get_scaling.clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = drk_render_func(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling + max_S, _ = torch.max(scales, dim=-1) + min_S, _ = torch.min(scales, dim=-1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + metrics["acutance_mean"] = float(torch.mean(self.gaussians.get_acutance)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = self.gaussians.get_scaling + sigma_sq = scales.max(dim=-1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_eagles.py b/methods/wrapper_eagles.py new file mode 100644 index 0000000000000000000000000000000000000000..d828768660a1a181121d407ca683de7921015c91 --- /dev/null +++ b/methods/wrapper_eagles.py @@ -0,0 +1,275 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from collections import OrderedDict + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../EAGLES_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModelSQ +from arguments import ModelParams, PipelineParams, OptimizationParams, QuantizeParams +from compress.inf_loss import EntropyLoss +from utils.general_utils import DecayScheduler + +@register_method("eagles") +class EAGLESWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser, {}) + self.op = OptimizationParams(self.parser, {}) + self.pp = PipelineParams(self.parser, {}) + self.qp = QuantizeParams(self.parser, {}) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.quantize = self.qp.extract(self.args) + + self.quantize.use_shift = [bool(el) for el in self.quantize.use_shift] + self.quantize.use_gumbel = [bool(el) for el in self.quantize.use_gumbel] + self.quantize.gumbel_period = [bool(el) for el in self.quantize.gumbel_period] + + self.gaussians = GaussianModelSQ(self.dataset.sh_degree, self.quantize) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + self.distortion_loss = EntropyLoss(self.gaussians.prob_models, lambdas=self.gaussians.ent_lambdas, noise_freq=self.quantize.noise_freq) + + self.temperature_scheds = OrderedDict() + for i, param in enumerate(self.gaussians.param_names): + self.temperature_scheds[param] = DecayScheduler( + total_steps=self.opt.iterations + 1, + decay_name='exp', + start=1.0, + end=self.quantize.temperature[i], + params={'temperature': self.quantize.temperature[i], 'decay_period': self.quantize.gumbel_period[i]} + ) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step, self.quantize) + + for i, param in enumerate(self.gaussians.param_names): + self.gaussians.latent_decoders[param].temperature = self.temperature_scheds[param](step) + self.gaussians.latent_decoders[param].use_gumbel = ((step / self.opt.iterations) < self.quantize.gumbel_period[i]) and self.quantize.use_gumbel[i] + + if (step - 1) % 10 == 0: + for i, param in enumerate(self.gaussians.param_names): + if hasattr(self.gaussians.latent_decoders[param], 'normalize'): + self.gaussians.latent_decoders[param].normalize(self.gaussians._latents[param]) + + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, get_infl=step < self.opt.prune_until_iter) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + loss_parasitic, bit_estimation = self.distortion_loss.loss(self.gaussians._latents, step, is_val=(self.quantize.noise_freq == 0)) + + if not isinstance(loss_parasitic, torch.Tensor): loss_parasitic = torch.tensor(float(loss_parasitic), device="cuda", requires_grad=True) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + g_t = self.gaussians._latents["xyz"].grad.clone() if self.gaussians._latents["xyz"].grad is not None else torch.zeros_like(self.gaussians._latents["xyz"]) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + g_p = self.gaussians._latents["xyz"].grad.clone() if self.gaussians._latents["xyz"].grad is not None else torch.zeros_like(self.gaussians._latents["xyz"]) + + state = self.gaussians.optimizer.state.get(self.gaussians._latents["xyz"], None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + for group in self.gaussians.optimizer.param_groups: + if group["name"] == "xyz": + eta_t = group['lr'] + break + else: + eta_t = 0.00016 + + eps0 = 1e-15 + u_t = (eta_t / (torch.sqrt(v_t) + eps0)) * g_t + u_p = (eta_t / (torch.sqrt(v_t) + eps0)) * g_p + + norm_p = torch.norm(u_p, dim=1) + valid_mask = (torch.norm(u_t, dim=1) > 0) & (norm_p > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(norm_p.mean() / (torch.norm(u_t, dim=1).mean() + norm_p.mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.prune_until_iter: + self.gaussians.add_influence_stats(render_pkg["influence"]) + + if step < self.opt.densify_until_iter: + if step <= self.opt.densify_from_iter or (step > self.opt.densify_from_iter and ((step - self.opt.densify_from_iter) % self.opt.densification_interval >= self.opt.accumulate_fraction * (self.opt.densification_interval - 1))): + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % self.opt.infl_prune_interval == 0 and step < self.opt.prune_until_iter: + self.gaussians.prune_influence(quantile_threshold=self.opt.quantile_threshold) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "quantization_bits": float(bit_estimation) if isinstance(bit_estimation, (float, torch.Tensor)) else 0.0 + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save_compressed(step, self.quantize) + + def load(self, model_path, iteration): + self.gaussians.load_compressed_pkl(os.path.join(model_path, f'iteration_{iteration}', 'compressed.pkl')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_edgeloss.py b/methods/wrapper_edgeloss.py new file mode 100644 index 0000000000000000000000000000000000000000..ad70b4cf2ed5e89ee9c8285879efefbfe7dfbc98 --- /dev/null +++ b/methods/wrapper_edgeloss.py @@ -0,0 +1,325 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +def _sobel(img): + """Single-scale luma Sobel response. NO stop-gradient — keeps autograd flowing.""" + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +@register_method("edgeloss") +class EdgeLossWrapper(BaseMethod): + """ + Ablation baseline answering Reviewer vm48: + "direct edge-weighted photometric losses or alternative + gradient-domain weighting strategies" + + Key contrast with SGF: NO stop-gradient on the weight map. + The Sobel weight is part of the autograd graph, so its gradient + flows back through all parameters (position, color, opacity, scale, rot), + not just positional gradients. This is the "trivial" edge-weighted + photometric loss SGF distinguishes itself from. + + Uses vanilla 3DGS backbone (3dgsAtlas_official) and standard + densification logic — only the loss is modified. + """ + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === Edge-weighted loss (NO stop-gradient — key contrast with SGF) === + edge_pred = _sobel(image) + edge_gt = _sobel(gt_image) + E = torch.abs(edge_pred - edge_gt).squeeze() + # Normalize without detaching: weight map participates in autograd. + W_edge = E / (E.max() + 1e-8) + + # L1 component, spatially reweighted by edge map (no detach) + L1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1 = ((1.0 - self.opt.lambda_dssim) * L1_map * (1.0 + 0.7 * W_edge)).mean() + + # SSIM component: standard scalar ssim (no spatial reweighting), + # keeping the comparison simple and the gradient through W_edge isolated to L1. + loss_ssim = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + loss_target = loss_l1 + loss_parasitic = loss_ssim + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone densification: 2-arg + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone signature with radii + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "edge_weight_mean": float(W_edge.mean()), + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_erankgs.py b/methods/wrapper_erankgs.py new file mode 100644 index 0000000000000000000000000000000000000000..485012d916b15756519be3ee286e98fca1b7b560 --- /dev/null +++ b/methods/wrapper_erankgs.py @@ -0,0 +1,265 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../erank_gs'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from utils.get_stats import get_effective_rank, get_ordered_scale_multiple + +@register_method("erankgs") +class ERankGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.erank_from_iter = 7000 + self.opt.no_new_densification = False + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, return_normal=True, return_depth=True) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + erank_loss = torch.tensor(0.0, device="cuda") + thin_loss = torch.tensor(0.0, device="cuda") + + scales = self.gaussians.get_scaling + erank = get_effective_rank(scales) + + if self.opt.erank_from_iter <= step <= self.opt.erank_end_iter: + erank_loss = self.opt.erank_lambda * torch.clamp(-torch.log(erank - 1 + 1e-7), 0).mean() + _, ordered_scale = get_ordered_scale_multiple(scales) + thin_loss = self.opt.thin_lambda * ordered_scale[:, 2].mean() + + loss_parasitic = erank_loss + thin_loss + loss = loss_target + loss_parasitic + + image_mask = viewpoint_cam.image_mask + if image_mask is not None: + o = render_pkg['render_opacity'] + opacity_mask_loss = -(image_mask * torch.log(o + 1e-7) + (1 - image_mask) * torch.log(1 - o + 1e-7)).mean() + loss += opacity_mask_loss + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and loss_parasitic.item() > 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._scaling.grad.clone() if self.gaussians._scaling.grad is not None else torch.zeros_like(self.gaussians._scaling) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._scaling.grad.clone() if self.gaussians._scaling.grad is not None else torch.zeros_like(self.gaussians._scaling) + + lr_scaling = 0.005 + v_t = torch.zeros_like(self.gaussians._scaling) + for group in self.gaussians.optimizer.param_groups: + if group["name"] == "scaling": + lr_scaling = group["lr"] + state = self.gaussians.optimizer.state.get(group['params'][0], None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"].clone() + break + + u_t = (lr_scaling / (torch.sqrt(v_t) + 1e-15)) * grad_target + u_p = (lr_scaling / (torch.sqrt(v_t) + 1e-15)) * grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + if self.opt.no_new_densification: + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + else: + self.gaussians.densify_and_prune_v2(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + metrics = { + "loss": float(loss), + "loss_l1": float(Ll1), + "loss_ssim": float(ssim_value), + "loss_erank": float(erank_loss), + "erank_mean": float(erank.mean()), + "needle_ratio": float((erank < 1.04).float().mean()), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + current_scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = current_scales + scales_2d = current_scales[:, :2] if current_scales.shape[1] >= 2 else current_scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, return_depth=True) + return {"image": render_pkg["render"], "depth": render_pkg.get("render_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_fds.py b/methods/wrapper_fds.py new file mode 100644 index 0000000000000000000000000000000000000000..0676eb5adf16c9417ae9dc9dd9124cedc7f641a5 --- /dev/null +++ b/methods/wrapper_fds.py @@ -0,0 +1,261 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '/root/autodl-tmp/FDS-GS_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.distribution_control import DistributionControl +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("fds") +class FDSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + self.distribution_control = DistributionControl(self.gaussians, self.scene.getTrainCameras().copy(), self.opt, self.pipe, self.dataset) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + v_stack = self.scene.getTrainCameras().copy(); self.viewpoint_stack = list(v_stack.values()) if isinstance(v_stack, dict) else v_stack + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + v_stack = self.scene.getTrainCameras().copy(); self.viewpoint_stack = list(v_stack.values()) if isinstance(v_stack, dict) else v_stack + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + metrics = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = [p.grad.clone() if p.grad is not None else torch.zeros_like(p) for p in params] + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = [p.grad.clone() if p.grad is not None else torch.zeros_like(p) for p in params] + + for group_name, params in param_groups_map.items(): + u_t_list = [] + u_p_list = [] + for i, p in enumerate(params): + state = self.gaussians.optimizer.state.get(p, {}) + v_t = state.get("exp_avg_sq", torch.zeros_like(p)) + lr = 1.0 + for pg in self.gaussians.optimizer.param_groups: + if pg["params"][0] is p: + lr = pg["lr"] + break + denom = torch.sqrt(v_t) + 1e-8 + u_t = (lr / denom) * grads_target[group_name][i] + u_p = (lr / denom) * grads_parasitic[group_name][i] + u_t_list.append(u_t.reshape(-1)) + u_p_list.append(u_p.reshape(-1)) + + M_Ut = torch.cat(u_t_list) + M_Up = torch.cat(u_p_list) + + if M_Ut.norm() > 0 and M_Up.norm() > 0: + cos = float(F.cosine_similarity(M_Ut.unsqueeze(0), M_Up.unsqueeze(0))) + r = float(M_Up.norm() / (M_Ut.norm() + M_Up.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + metrics[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.distribution_control.control(step) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics.update({ + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + }) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_fdsgs.py b/methods/wrapper_fdsgs.py new file mode 100644 index 0000000000000000000000000000000000000000..0dcf462178547b39c0907de2d7812f5d3d0e89d0 --- /dev/null +++ b/methods/wrapper_fdsgs.py @@ -0,0 +1,267 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../FDS-GS_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.distribution_control import DistributionControl +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("fdsgs") +class FDSGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.white_background = False + self.args.random_background = False + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + + self.distribution_control = DistributionControl(self.gaussians, self.scene.getTrainCameras().copy(), self.opt, self.pipe, self.dataset) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + self.group_mapping = { + "spatial": ["xyz"], + "geometry": ["scaling", "rotation", "R"], + "opacity": ["opacity"], + "appearance": ["f_dc", "f_rest"] + } + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + keys = list(self.viewpoint_stack.keys()) + viewpoint_cam = self.viewpoint_stack.pop(keys[random.randint(0, len(keys) - 1)]) + + bg = torch.rand(3, device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_loss = ssim(image, gt_image) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_loss) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step > 0 and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + lr_dict = {group['name']: group['lr'] for group in self.gaussians.optimizer.param_groups} + param_dict = {group['name']: group['params'][0] for group in self.gaussians.optimizer.param_groups} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + u_targets = {} + for g_name, p_names in self.group_mapping.items(): + u_list = [] + for p_name in p_names: + p = param_dict.get(p_name) + if p is not None and p.grad is not None: + state = self.gaussians.optimizer.state.get(p) + v = state.get('exp_avg_sq', torch.zeros_like(p)) if state else torch.zeros_like(p) + lr = lr_dict.get(p_name, 1e-3) + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + u_list.append(u.reshape(-1)) + if u_list: + u_targets[g_name] = torch.cat(u_list) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + u_parasitics = {} + for g_name, p_names in self.group_mapping.items(): + u_list = [] + for p_name in p_names: + p = param_dict.get(p_name) + if p is not None and p.grad is not None: + state = self.gaussians.optimizer.state.get(p) + v = state.get('exp_avg_sq', torch.zeros_like(p)) if state else torch.zeros_like(p) + lr = lr_dict.get(p_name, 1e-3) + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + u_list.append(u.reshape(-1)) + if u_list: + u_parasitics[g_name] = torch.cat(u_list) + + for g_name in self.group_mapping.keys(): + ut = u_targets.get(g_name) + up = u_parasitics.get(g_name) + if ut is not None and up is not None and ut.norm() > 0 and up.norm() > 0: + cos = float(F.cosine_similarity(ut.unsqueeze(0), up.unsqueeze(0))) + r = float(up.norm() / (ut.norm() + up.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{g_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.distribution_control.control(step) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step > 0 and step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"]} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.gaussians._scaling) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_flod.py b/methods/wrapper_flod.py new file mode 100644 index 0000000000000000000000000000000000000000..47a4ede13a58dc20b94e13499984f38287085128 --- /dev/null +++ b/methods/wrapper_flod.py @@ -0,0 +1,329 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/FLoD_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def expand_list(l, size): + if len(l) >= size: + return l[:size] + return l + [l[-1]] * (size - len(l)) + +@register_method("flod") +class FLODWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel( + sh_degree=self.dataset.sh_degree, + lod1_scaling_lower_bound=self.dataset.lod1_scaling_lower_bound, + lod_scaling_ratio=self.dataset.lod_scaling_ratio, + increase_lod_num_childs=self.dataset.increase_lod_num_childs, + current_lod=self.opt.lod_min, + max_lod=self.opt.lod_max, + use_voxel_sampling=self.dataset.use_voxel_sampling, + voxel_sampling_size=self.dataset.voxel_sampling_size + ) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.current_lod_idx = 0 + + lods = self.opt.lod_max - self.opt.lod_min + 1 + self.densify_grad_thresholds = expand_list(self.opt.densify_grad_threshold, lods) + self.densification_intervals = expand_list(self.opt.densification_interval, lods) + self.densify_from_iters = expand_list(self.opt.densify_from_iter, lods) + self.densify_until_iters = expand_list(self.opt.densify_until_iter, lods) + self.prune_opacity_thresholds = expand_list(self.opt.prune_opacity_threshold, lods) + self.pruning_intervals = expand_list(self.opt.pruning_interval, lods) + self.prune_from_iters = expand_list(self.opt.prune_from_iter, lods) + self.prune_until_iters = expand_list(self.opt.prune_until_iter, lods) + self.prune_overlap_thresholds = expand_list(self.opt.prune_overlap_threshold, lods) + self.pruning_overlap_intervals = expand_list(self.opt.pruning_overlap_interval, lods) + self.opacity_reset_intervals = expand_list(self.opt.opacity_reset_interval, lods) + self.lambda_dssims = expand_list(self.opt.lambda_dssim, lods) + + def train_iteration(self, step): + target_idx = 0 + local_step = step + for idx, iters in enumerate(self.opt.lod_iterations): + if local_step <= iters: + target_idx = idx + break + local_step -= iters + + if target_idx > self.current_lod_idx: + self.gaussians.increase_lod() + self.gaussians.training_setup(self.opt) + self.current_lod_idx = target_idx + + self.gaussians.update_learning_rate(local_step) + + if local_step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + lambda_dssim = self.lambda_dssims[self.current_lod_idx] + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_target = (1.0 - lambda_dssim) * Ll1 + loss_parasitic = lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and local_step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + def get_effective_steps(loss_comp): + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_comp.backward(retain_graph=True) + grads_eff = {} + for group_name, params in param_groups_map.items(): + grad_list = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v_t = state.get("exp_avg_sq", torch.zeros_like(p)) + group_lr = 0.0 + for pg in self.gaussians.optimizer.param_groups: + if id(p) in [id(pgp) for pgp in pg['params']]: + group_lr = pg['lr'] + break + u = (group_lr / (torch.sqrt(v_t) + 1e-8)) * p.grad.clone() + grad_list.append(u.reshape(-1)) + if grad_list: + grads_eff[group_name] = torch.cat(grad_list) + return grads_eff + + grads_target_eff = get_effective_steps(loss_target) + grads_parasitic_eff = get_effective_steps(loss_parasitic) + + for group_name in param_groups_map: + gt = grads_target_eff.get(group_name) + gp = grads_parasitic_eff.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + densify_grad_threshold = self.densify_grad_thresholds[self.current_lod_idx] + densification_interval = self.densification_intervals[self.current_lod_idx] + densify_from_iter = self.densify_from_iters[self.current_lod_idx] + densify_until_iter = self.densify_until_iters[self.current_lod_idx] + prune_opacity_threshold = self.prune_opacity_thresholds[self.current_lod_idx] + pruning_interval = self.pruning_intervals[self.current_lod_idx] + prune_from_iter = self.prune_from_iters[self.current_lod_idx] + prune_until_iter = self.prune_until_iters[self.current_lod_idx] + prune_overlap_threshold = self.prune_overlap_thresholds[self.current_lod_idx] + pruning_overlap_interval = self.pruning_overlap_intervals[self.current_lod_idx] + opacity_reset_interval = self.opacity_reset_intervals[self.current_lod_idx] + + if local_step < densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if local_step > densify_from_iter and local_step % densification_interval == 0: + self.gaussians.densify(densify_grad_threshold, extent=self.scene.cameras_extent) + + if local_step < prune_until_iter: + if local_step > prune_from_iter and local_step % pruning_interval == 0: + size_threshold = 20 if (local_step > opacity_reset_interval and (self.opt.lod_min + self.current_lod_idx) >= self.opt.prune_max_radii2D_lod) else None + + if local_step % pruning_overlap_interval == 0: + self.gaussians.prune(prune_opacity_threshold, self.scene.cameras_extent, size_threshold, prune_overlap_threshold) + else: + self.gaussians.prune(prune_opacity_threshold, self.scene.cameras_extent, size_threshold, 0.0) + + if local_step < prune_until_iter: + if local_step % opacity_reset_interval == 0 or (self.dataset.white_background and local_step == densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "current_lod": int(self.opt.lod_min + self.current_lod_idx), + "local_step": int(local_step) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = self.gaussians.get_scaling.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step, self.opt.lod_min + self.current_lod_idx) + + def load(self, model_path, iteration): + self.gaussians.load(model_path, load_iteration=iteration, load_independent_lvl=False) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = self.gaussians.get_rotation.clone() + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = self.gaussians.get_scaling + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gaussian_surfel.py b/methods/wrapper_gaussian_surfel.py new file mode 100644 index 0000000000000000000000000000000000000000..48450990c41195dd9d97cdf19f6b841aa4444082 --- /dev/null +++ b/methods/wrapper_gaussian_surfel.py @@ -0,0 +1,189 @@ +import os +import sys +import torch +import random +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +SURFEL_ROOT = "/root/autodl-tmp/gaussian_surfels" +if SURFEL_ROOT not in sys.path: + sys.path.insert(0, SURFEL_ROOT) + +@register_method("gaussian_surfel") +class SurfelWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + from scene import Scene, GaussianModel + from arguments import ModelParams, PipelineParams, OptimizationParams + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.resolution = dataset_config.get("resolution", 1) + self.args.eval = True + self.args.surface = True + if any(x in self.args.source_path.lower() for x in ["tnt", "360", "tanks"]): + self.args.images = "images" + else: + self.args.images = "image" + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + _prev_cwd = os.getcwd(); os.chdir("/root/autodl-tmp/gaussian_surfels") + self.gaussians = GaussianModel(self.dataset) + os.chdir(_prev_cwd) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians, self.opt.camera_lr, shuffle=False, resolution_scales=[1, 2, 4]) + self.gaussians.training_setup(self.opt) + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = None + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.cap_gaussians = hyperparams.get("cap_gaussians", None) + + def train_iteration(self, step): + from gaussian_renderer import render as native_render + from utils.loss_utils import l1_loss, ssim, cos_loss + from utils.image_utils import depth2normal + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + scale = 1 + if step < 2000: + scale = 4 + elif step < 5000: + scale = 2 + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras(scale).copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, [float('inf'), float('inf')]) + image, normal, depth, opac = render_pkg["render"], render_pkg["normal"], render_pkg["depth"], render_pkg["opac"] + gt_image = viewpoint_cam.get_gtImage(self.background, self.dataset.use_mask) + mask_vis = (opac.detach() > 1e-5) + loss_target = (1.0 - self.opt.lambda_dssim) * l1_loss(image, gt_image) + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + d2n = depth2normal(depth, mask_vis, viewpoint_cam) + normal_norm = F.normalize(normal, dim=0) * mask_vis + loss_surface = cos_loss(normal_norm, d2n) + loss = loss_target + loss_parasitic + loss += (0.01 + 0.1 * min(2 * step / self.opt.iterations, 1)) * loss_surface + grad_cos_sim = 0.0 + p_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad() + loss_target.backward(retain_graph=True) + g_t = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else None + self.gaussians.optimizer.zero_grad() + loss_parasitic.backward(retain_graph=True) + g_p = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else None + if g_t is not None and g_p is not None: + mask = (g_t.norm(dim=1) > 1e-8) & (g_p.norm(dim=1) > 1e-8) + if mask.any(): + grad_cos_sim = float(F.cosine_similarity(g_t[mask], g_p[mask], dim=1).mean()) + p_ratio = float(g_p[mask].norm(dim=1).mean() / (g_t[mask].norm(dim=1).mean() + 1e-8)) + self.gaussians.optimizer.zero_grad() + loss.backward() + else: + loss.backward() + + # 修复点：包裹在 no_grad() 中，且在 optimizer.step() 之前执行 + with torch.no_grad(): + if step > self.opt.densify_from_iter: + self.gaussians.max_radii2D[render_pkg["visibility_filter"]] = torch.max(self.gaussians.max_radii2D[render_pkg["visibility_filter"]], render_pkg["radii"][render_pkg["visibility_filter"]]) + self.gaussians.add_densification_stats(render_pkg["viewspace_points"], render_pkg["visibility_filter"]) + if step % self.opt.densification_interval == 0: + self.gaussians.adaptive_prune(0.1, self.scene.cameras_extent) + if self.cap_gaussians is None or len(self.gaussians.get_xyz) < self.cap_gaussians: + self.gaussians.adaptive_densify(self.opt.densify_grad_threshold, self.scene.cameras_extent) + if step % self.opt.opacity_reset_interval == 0: + self.gaussians.reset_opacity(0.12, step) + + # 统计完致密化状态后，再进行优化更新并清空梯度 + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = len(self.gaussians.get_xyz) + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": num_gaussians, + "delta_N": num_gaussians - self.last_n_gaussians, + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024**3)), + "grad_cos_sim": grad_cos_sim, + "parasitic_ratio": p_ratio + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).detach() + scales = torch.exp(self.gaussians._scaling).detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + return metrics, histograms + + def render(self, camera): + from gaussian_renderer import render as native_render + with torch.no_grad(): + pkg = native_render(camera, self.gaussians, self.pipe, self.background, [float('inf'), float('inf')]) + return {"image": pkg["render"], "depth": pkg["depth"]} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, f"point_cloud/iteration_{iteration}/point_cloud.ply")) + + def get_spatial_centers(self): + return self.gaussians._xyz.detach() + + def compute_physical_metrics(self, cameras=None): + scales = self.gaussians.get_scaling.detach() + return { + "gamma_median": float(torch.median(scales[:, 0] / (scales[:, 1] + 1e-7))), + "alpha_mean": float(self.gaussians.get_opacity.detach().mean()), + "z_scale_mean": float(scales[:, 2].mean()) + } + + def evaluate_spatial_field(self, query_points, cameras=None): + from pytorch3d.ops import knn_points + xyz = self.gaussians.get_xyz.detach() + opac = self.gaussians.get_opacity.detach() + dist_sq, _, _ = knn_points(query_points.unsqueeze(0), xyz.unsqueeze(0), K=1) + dist = torch.sqrt(dist_sq.squeeze(0)) + weights = torch.exp(-0.5 * (dist / 0.01)**2) + return (weights * opac.T).sum(dim=1) diff --git a/methods/wrapper_gaussianfocus.py b/methods/wrapper_gaussianfocus.py new file mode 100644 index 0000000000000000000000000000000000000000..2a23e42aa600fd7b43c4527df967146f725077af --- /dev/null +++ b/methods/wrapper_gaussianfocus.py @@ -0,0 +1,373 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn as nn +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GaussianFocus_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +class PatchAttention(nn.Module): + def __init__(self, channel): + super(PatchAttention, self).__init__() + self.query_conv = nn.Conv2d(channel, 1, kernel_size=1) + self.key_conv = nn.Conv2d(channel, 1, kernel_size=1) + self.value_conv = nn.Conv2d(channel, channel, kernel_size=1) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, image, gt_image, block_size=64): + batch_size, C, height, width = image.size() + blocks = [] + for i in range(0, height, block_size): + for j in range(0, width, block_size): + block_image = image[:, :, i:i+block_size, j:j+block_size] + block_gt_image = gt_image[:, :, i:i+block_size, j:j+block_size] + + pad_height = block_size - block_image.size(2) + pad_width = block_size - block_image.size(3) + if pad_height > 0 or pad_width > 0: + block_image = F.pad(block_image, (0, pad_width, 0, pad_height)) + block_gt_image = F.pad(block_gt_image, (0, pad_width, 0, pad_height)) + + query = self.query_conv(block_image).view(batch_size, -1, block_size * block_size).permute(0, 2, 1) + key = self.key_conv(block_gt_image).view(batch_size, -1, block_size * block_size) + value = self.value_conv(block_gt_image).view(batch_size, -1, block_size * block_size) + + attention = self.softmax(torch.bmm(query, key)) + out = torch.bmm(value, attention.permute(0, 2, 1)) + out = out.view(batch_size, C, block_size, block_size) + blocks.append(out[:, :, :min(block_size, height - i), :min(block_size, width - j)]) + + output = torch.cat([torch.cat(blocks[i * (width // block_size + (width % block_size > 0)):(i + 1) * (width // block_size + (width % block_size > 0))], dim=3) for i in range(height // block_size + (height % block_size > 0))], dim=2) + return output + +def edge_loss(input_img, target_img): + edge_filter = torch.Tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]).view(1, 1, 3, 3).to(input_img.device) + edge_filter = edge_filter.repeat(input_img.size(1), 1, 1, 1) + edge_filter = nn.Parameter(data=edge_filter, requires_grad=False) + + input_x = F.conv2d(input_img, edge_filter, padding=1, groups=input_img.size(1)) + input_y = F.conv2d(input_img, edge_filter.transpose(2, 3), padding=1, groups=input_img.size(1)) + target_x = F.conv2d(target_img, edge_filter, padding=1, groups=target_img.size(1)) + target_y = F.conv2d(target_img, edge_filter.transpose(2, 3), padding=1, groups=target_img.size(1)) + + loss_x = F.l1_loss(input_x, target_x) + loss_y = F.l1_loss(input_y, target_y) + return (loss_x + loss_y) / 2 + +def frequency_loss(input_img, target_img): + grad_x_input = input_img[:, :, 1:] - input_img[:, :, :-1] + grad_y_input = input_img[:, 1:, :] - input_img[:, :-1, :] + grad_x_target = target_img[:, :, 1:] - target_img[:, :, :-1] + grad_y_target = target_img[:, 1:, :] - target_img[:, :-1, :] + + loss_x = F.l1_loss(grad_x_input, grad_x_target) + loss_y = F.l1_loss(grad_y_input, grad_y_target) + return (loss_x + loss_y) / 2 + +class CombinedLoss(nn.Module): + def __init__(self): + super(CombinedLoss, self).__init__() + + def forward(self, finals, gt_image, beta=1, eta=1): + if len(finals.shape) == 3: + finals = finals.unsqueeze(0) + if len(gt_image.shape) == 3: + gt_image = gt_image.unsqueeze(0) + e_loss = edge_loss(finals, gt_image) + f_loss = frequency_loss(finals, gt_image) + return beta * e_loss + eta * f_loss + +@register_method("gaussianfocus") +class GaussianFocusWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.dataset.kernel_size = 0.1 + self.dataset.ray_jitter = False + self.dataset.resample_gt_image = False + self.dataset.sample_more_highres = False + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.gaussians.compute_3D_filter(cameras=self.scene.getTrainCameras()) + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + self.patch_attention = PatchAttention(channel=3).to("cuda") + self.attn_optimizer = torch.optim.Adam(self.patch_attention.parameters(), lr=1e-3) + self.combined_loss_fn = CombinedLoss() + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.gaussians.compute_3D_filter(cameras=self.scene.getTrainCameras()) + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + patch_attn_active = 0 + finals = image + if step % 50 == 0: + patch_attn_active = 1 + image1 = image.unsqueeze(0) + gt_image1 = gt_image.unsqueeze(0) + output = self.patch_attention(image1, gt_image1) + finals = output * image + + finals = torch.squeeze(finals, 0) + finals_var = finals + gt_image_var = gt_image + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + loss_x = self.combined_loss_fn(finals_var, gt_image_var) + loss = loss_target + loss_x + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = { + "sti_spatial": 0.0, + "sti_geometry": 0.0, + "sti_opacity": 0.0, + "sti_appearance": 0.0 + } + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if getattr(loss_x, "requires_grad", False): loss_x.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if getattr(loss_x, "requires_grad", False): loss_x.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + N_total = self.gaussians.get_xyz.shape[0] + SAFE_N = 2_000_000 # above this, skip group-wise STI to avoid OOM + for group_name in param_groups_map: + if N_total > SAFE_N: + stats[f"sti_{group_name}"] = 0.0 + continue + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + # 1D dot product is memory-cheap; cosine_similarity on unsqueeze(0) + # allocates 2D intermediates that OOM at N>2M. + cos = float((gt @ gp) / (gt.norm() * gp.norm() + 1e-7)) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.attn_optimizer.zero_grad() + loss.backward() + else: + self.attn_optimizer.zero_grad() + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 0.02 if step > 1000 else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, step) + self.gaussians.compute_3D_filter(cameras=self.scene.getTrainCameras()) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter: + if step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.scene.getTrainCameras()) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + if patch_attn_active: + self.attn_optimizer.step() + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_x), + "loss_edge_freq": float(loss_x), + "patch_attn_active": int(patch_attn_active), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gaussianpro.py b/methods/wrapper_gaussianpro.py new file mode 100644 index 0000000000000000000000000000000000000000..5255f446462a0811f8d3cbec63a9204127f5303a --- /dev/null +++ b/methods/wrapper_gaussianpro.py @@ -0,0 +1,273 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../gaussianpro_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("gaussianpro") +class GaussianProWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.normal_loss = True + self.opt.flatten_loss = True + self.opt.sparse_loss = True + self.opt.lambda_flatten = 100.0 + self.opt.lambda_l1_normal = 0.01 + self.opt.lambda_cos_normal = 0.01 + self.opt.lambda_sparse = 0.001 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, return_normal=self.opt.normal_loss, return_opacity=True) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + opacity_mask = render_pkg["render_opacity"] > 0.0 + opacity_mask = opacity_mask.unsqueeze(0).expand(3, -1, -1) + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image[opacity_mask], gt_image[opacity_mask]) + ssim_value = ssim(image, gt_image, mask=opacity_mask) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_target_ssim = self.opt.lambda_dssim * (1.0 - ssim_value) + loss_target_total = loss_target + loss_target_ssim + + loss_parasitic = torch.tensor(0.0, device="cuda") + loss_flatten_val = 0.0 + loss_normal_val = 0.0 + + if self.opt.flatten_loss: + scales = self.gaussians.get_scaling + min_scale, _ = torch.min(scales, dim=1) + min_scale = torch.clamp(min_scale, 0, 30) + flatten_loss_tensor = torch.abs(min_scale).mean() + loss_parasitic = loss_parasitic + self.opt.lambda_flatten * flatten_loss_tensor + loss_flatten_val = float(flatten_loss_tensor) + + if self.opt.sparse_loss: + opacity = self.gaussians.get_opacity.clamp(1e-6, 1-1e-6) + log_opacity = opacity * torch.log(opacity) + log_one_minus_opacity = (1 - opacity) * torch.log(1 - opacity) + sparse_loss_tensor = -1 * (log_opacity + log_one_minus_opacity)[visibility_filter].mean() + loss_parasitic = loss_parasitic + self.opt.lambda_sparse * sparse_loss_tensor + + if self.opt.normal_loss and viewpoint_cam.normal is not None: + rendered_normal = render_pkg["render_normal"] + normal_gt = viewpoint_cam.normal.cuda() + filter_mask = (normal_gt != -10)[0, :, :].to(torch.bool) + if filter_mask.any(): + l1_normal = torch.abs(rendered_normal - normal_gt).sum(dim=0)[filter_mask].mean() + cos_normal = (1.0 - torch.sum(rendered_normal * normal_gt, dim=0))[filter_mask].mean() + normal_loss_tensor = self.opt.lambda_l1_normal * l1_normal + self.opt.lambda_cos_normal * cos_normal + loss_parasitic = loss_parasitic + normal_loss_tensor + loss_normal_val = float(normal_loss_tensor) + + loss = loss_target_total + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and loss_parasitic.requires_grad: + track_params = [ + self.gaussians._xyz, + self.gaussians._scaling, + self.gaussians._opacity, + self.gaussians._rotation + ] + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target_total.backward(retain_graph=True) + grad_t_list = [p.grad.clone().view(p.shape[0], -1) if p.grad is not None else torch.zeros_like(p).view(p.shape[0], -1) for p in track_params] + grad_target_cat = torch.cat(grad_t_list, dim=1) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_p_list = [p.grad.clone().view(p.shape[0], -1) if p.grad is not None else torch.zeros_like(p).view(p.shape[0], -1) for p in track_params] + grad_parasitic_cat = torch.cat(grad_p_list, dim=1) + + valid_mask = (torch.norm(grad_target_cat, dim=1) > 0) & (torch.norm(grad_parasitic_cat, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target_cat[valid_mask], grad_parasitic_cat[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic_cat, dim=1).mean() / (torch.norm(grad_target_cat, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_target_ssim), + "loss_flatten": loss_flatten_val, + "loss_normal": loss_normal_val, + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + max_S, _ = torch.max(scales, dim=-1) + min_S, _ = torch.min(scales, dim=-1) + histograms["anisotropy"] = max_S / (min_S + 1e-7) + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("render_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, "point_cloud", f"iteration_{iteration}", "point_cloud.ply")) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.gaussians._scaling) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gaussianshader.py b/methods/wrapper_gaussianshader.py new file mode 100644 index 0000000000000000000000000000000000000000..0a3a049fcb627301a037618199d6570163cc1c8b --- /dev/null +++ b/methods/wrapper_gaussianshader.py @@ -0,0 +1,239 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GaussianShader'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams, SGFParams +from utils.sgf_utils import SGFController + +@register_method("gaussianshader") +class GaussianShaderWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.sp = SGFParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.sgf_args = self.sp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + self.sgf_controller = SGFController(self.sgf_args) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.train_cameras = self.scene.getTrainCameras().copy() + self.viewpoint_stack = self.train_cameras.copy() + + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.train_cameras.copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + loss, struct_map, l1_map = self.sgf_controller.compute_loss(image, gt_image, lambda_dssim=self.opt.lambda_dssim) + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_ssim_val = 1.0 - ssim_value + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + Ll1.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic = loss - Ll1 + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + eta_t = self.opt.position_lr_init + u_target = (eta_t / (torch.sqrt(v_t) + 1e-15)) * grad_target + u_parasitic = (eta_t / (torch.sqrt(v_t) + 1e-15)) * grad_parasitic + + valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter: + if step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(Ll1), "loss_ssim": float(loss_ssim_val), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_struct_map_mean": float(struct_map.mean()), + "sgf_l1_map_mean": float(l1_map.mean()), + "filter_3D_mean": float(self.gaussians.filter_3D.mean()) if hasattr(self.gaussians, "filter_3D") else 0.0 + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + opacities = self.gaussians.get_opacity_with_3D_filter.clone().detach() + scales = self.gaussians.get_scaling_with_3D_filter.clone().detach() + histograms["opacity"] = opacities + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling_with_3D_filter + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.gaussians.get_opacity_with_3D_filter)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = self.gaussians.get_opacity_with_3D_filter.squeeze() + scales = self.gaussians.get_scaling_with_3D_filter + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gaussianwild.py b/methods/wrapper_gaussianwild.py new file mode 100644 index 0000000000000000000000000000000000000000..b1ab53c31311cf05112bdb563185752a71f7ca1f --- /dev/null +++ b/methods/wrapper_gaussianwild.py @@ -0,0 +1,276 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Gaussian-Wild_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams, args_init +import lpips + +@register_method("gaussianwild") +class GaussianWildWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args, _ = self.parser.parse_known_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config.get("model_path", "") + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.scene_name = "benchmark" + self.args.data_perturb = [] + self.args.use_colors_precomp = True + self.args.use_decode_with_pos = False + self.args.use_indep_mask_branch = False + self.args.use_features_mask = True + self.args.features_mask_loss_coef = 0.15 + self.args.features_mask_iters = 2500 + self.args.use_okmap = False + self.args.use_kmap_pjmap = True + self.args.map_num = 3 + self.args.use_wo_adative = 0 + self.args.use_xw_init_box_coord = True + self.args.use_color_net = True + self.args.use_scaling_loss = False + self.args.use_lpips_loss = True + self.args.use_box_coord_loss = True + self.args.coord_scale = 1 + self.args.device = "cuda:0" + + self.args = args_init.argument_init(self.args) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree, self.args) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians, shuffle=False) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.lpips_criteria = lpips.LPIPS(net='vgg').to("cuda:0") + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step, self.opt.warm_up_iter) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + loss_target = torch.tensor(0.0, device="cuda") + loss_parasitic = torch.tensor(0.0, device="cuda") + + if self.args.use_features_mask and step > self.args.features_mask_iters: + mask = self.gaussians.features_mask + mask = F.interpolate(mask, size=(image.shape[-2:])) + Ll1 = l1_loss(image * mask, gt_image * mask) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image * mask, gt_image * mask)) + loss_parasitic += (torch.square(1 - self.gaussians.features_mask)).mean() * self.args.features_mask_loss_coef + else: + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + if self.args.use_scaling_loss: + loss_parasitic += torch.abs(self.gaussians.get_scaling).mean() * self.args.scaling_loss_coef + + if self.args.use_lpips_loss: + loss_target += self.lpips_criteria(image, gt_image).mean() * self.opt.lpips_loss_coef + + if (self.args.use_kmap_pjmap or self.args.use_okmap) and self.args.use_box_coord_loss: + if hasattr(self.gaussians, "map_pts_norm"): + loss_parasitic += torch.relu(torch.abs(self.gaussians.map_pts_norm) - 1).mean() * self.opt.box_coord_loss_coef + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + u_target = grad_target + u_parasitic = grad_parasitic + + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = self.opt.position_lr_init + for pg in self.gaussians.optimizer.param_groups: + if pg["name"] == "xyz": + lr = pg["lr"] + u_target = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_parasitic = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.opacity_threshold, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + mask_activation = float(self.gaussians.features_mask.mean()) if hasattr(self.gaussians, "features_mask") else 0.0 + coord_bound = float(torch.abs(self.gaussians.map_pts_norm).mean()) if hasattr(self.gaussians, "map_pts_norm") else 0.0 + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "mask_activation_mean": mask_activation, "box_coord_out_of_bound": coord_bound + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_intrinsic.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + pt_path = os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply') + self.gaussians.load_ckpt_ply(pt_path) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_geogaussian.py b/methods/wrapper_geogaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..943944c4fa181c915f027d761787c540638deb7f --- /dev/null +++ b/methods/wrapper_geogaussian.py @@ -0,0 +1,297 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GeoGaussian'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from diff_gaussian_rasterization import SurfaceAlign +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("geogaussian") +class GeoGaussianWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.KNN_index = None + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic_ssim = self.opt.lambda_dssim * (1.0 - ssim_value) + + pair_d_loss = torch.tensor(0.0, device="cuda") + pair_normal_loss = torch.tensor(0.0, device="cuda") + + if self.KNN_index is not None and visibility_filter.sum() > 0: + mask = torch.logical_and(visibility_filter, self.gaussians.get_type.squeeze() == 1) + if mask.sum() > 0: + pair_d_loss_raw, pair_normal_loss_raw = SurfaceAlign()( + self.gaussians.get_xyz, + self.gaussians.get_xyz_id.contiguous(), + self.gaussians.get_rotation, + self.KNN_index[mask] + ) + pair_d_loss = 0.05 * torch.mean(pair_d_loss_raw) + pair_normal_loss = 0.01 * torch.mean(pair_normal_loss_raw) + + loss_parasitic = loss_parasitic_ssim + pair_d_loss + pair_normal_loss + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group in self.gaussians.optimizer.param_groups: + name = group["name"] + p = group["params"][0] + if p.grad is not None: + grads_target[name] = p.grad.clone() + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group in self.gaussians.optimizer.param_groups: + name = group["name"] + p = group["params"][0] + if p.grad is not None: + grads_parasitic[name] = p.grad.clone() + + gt_xyz = grads_target.get("xyz", torch.zeros(1, device="cuda")) + gp_xyz = grads_parasitic.get("xyz", torch.zeros(1, device="cuda")) + valid_mask = (torch.norm(gt_xyz, dim=1) > 0) & (torch.norm(gp_xyz, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(gt_xyz[valid_mask], gp_xyz[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(gp_xyz, dim=1).mean() / (torch.norm(gt_xyz, dim=1).mean() + 1e-7)) + + group_mapping = { + "spatial": ["xyz"], + "geometry": ["scaling", "rotation"], + "opacity": ["opacity"], + "appearance": ["f_dc", "f_rest"] + } + + for super_group, sub_names in group_mapping.items(): + u_t_list = [] + u_p_list = [] + for name in sub_names: + for group in self.gaussians.optimizer.param_groups: + if group["name"] == name: + p = group["params"][0] + state = self.gaussians.optimizer.state.get(p) + vt = state.get("exp_avg_sq", torch.zeros_like(p)) if state else torch.zeros_like(p) + lr = group["lr"] + gt = grads_target.get(name, torch.zeros_like(p)) + gp = grads_parasitic.get(name, torch.zeros_like(p)) + step_t = (lr / (torch.sqrt(vt) + 1e-15)) * gt + step_p = (lr / (torch.sqrt(vt) + 1e-15)) * gp + u_t_list.append(step_t.reshape(-1)) + u_p_list.append(step_p.reshape(-1)) + + if u_t_list and u_p_list: + U_T = torch.cat(u_t_list) + U_P = torch.cat(u_p_list) + if U_T.norm() > 0 and U_P.norm() > 0: + cos = float(F.cosine_similarity(U_T.unsqueeze(0), U_P.unsqueeze(0))) + r = float(U_P.norm() / (U_T.norm() + U_P.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{super_group}"] = ti + else: + stats[f"sti_{super_group}"] = 0.0 + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) + self.KNN_index = self.gaussians.findKNN() + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + elif step > self.opt.densify_until_iter: + if step % 3000 == 0: + self.KNN_index = self.gaussians.findKNN() + + if step <= self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic_ssim), + "loss_pair_d": float(pair_d_loss), "loss_pair_normal": float(pair_normal_loss), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("render_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_ges.py b/methods/wrapper_ges.py new file mode 100644 index 0000000000000000000000000000000000000000..13cc13d9e4c9227ee9f5fc1df4c8babe782ce04d --- /dev/null +++ b/methods/wrapper_ges.py @@ -0,0 +1,243 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../ges_official'))) +from utils.loss_utils import l1_loss, ssim +from utils.image_utils import apply_dog_filter +from gaussian_renderer import render_laplacian +from scene import Scene, LaplacianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("ges") +class GESWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = LaplacianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = render_laplacian(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + freq = (step / self.opt.iterations) * 100 + mask = apply_dog_filter(image.unsqueeze(0), freq=freq, scale_factor=self.opt.im_laplace_scale_factor).squeeze(0) + mask_loss = l1_loss(image * mask, gt_image * mask) + + loss_target = (1.0 - self.opt.lambda_dssim - self.opt.lambda_im_laplace) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + self.opt.lambda_im_laplace * mask_loss + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr_xyz = 0.0 + for param_group in self.gaussians.optimizer.param_groups: + if param_group["name"] == "xyz": + lr_xyz = param_group["lr"] + break + + u_target = (lr_xyz / (torch.sqrt(v_t) + 1e-15)) * grad_target + u_parasitic = (lr_xyz / (torch.sqrt(v_t) + 1e-15)) * grad_parasitic + + norm_target = torch.norm(u_target, dim=1) + norm_parasitic = torch.norm(u_parasitic, dim=1) + valid_mask = (norm_target > 0) & (norm_parasitic > 0) + + if valid_mask.any(): + cos_sim = F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1) + grad_cos_sim = float(cos_sim.mean()) + + r_group = norm_parasitic[valid_mask] / (norm_target[valid_mask] + norm_parasitic[valid_mask] + 1e-7) + ti = r_group * torch.clamp(-cos_sim, min=0.0) + parasitic_ratio = float(ti.mean()) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.prune_opacity_threshold, self.scene.cameras_extent, size_threshold) + + if step > self.opt.densify_from_iter and step % self.opt.shape_pruning_interval == 0: + self.gaussians.size_prune(self.opt.prune_shape_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % self.opt.shape_reset_interval == 0: + self.gaussians.reset_shape() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "average_shape": float(self.gaussians.get_shape.mean()) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = (torch.exp(self.gaussians._scaling) * self.gaussians.get_shape).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + histograms["shape"] = self.gaussians.get_shape.clone().detach() + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = render_laplacian(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.gaussians._scaling) * self.gaussians.get_shape + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + metrics["shape_mean"] = float(torch.mean(self.gaussians.get_shape)) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points, cameras=None): + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) * self.gaussians.get_shape + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_ghap.py b/methods/wrapper_ghap.py new file mode 100644 index 0000000000000000000000000000000000000000..f649672d24b5f953776fbe883c37e616d32745fe --- /dev/null +++ b/methods/wrapper_ghap.py @@ -0,0 +1,261 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../ghap_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from train_and_prune import subsampling + +@register_method("ghap") +class GHAPWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.compaction_iter = hyperparams.get("sampling_iter", 15000) + self.compaction_ratio = hyperparams.get("sampling_ratio", 0.1) + self.gmr_trigger_status = 0.0 + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree, self.opt.optimizer_type) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step, False) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, use_trained_exp=self.dataset.train_test_exp, separate_sh=(self.opt.optimizer_type == "sparse_adam")) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + if viewpoint_cam.alpha_mask is not None: + image *= viewpoint_cam.alpha_mask.cuda() + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + if hasattr(self.gaussians, "exposure_optimizer"): + self.gaussians.exposure_optimizer.zero_grad(set_to_none=True) + + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + p_xyz = None + lr_xyz = 0.0 + for group in self.gaussians.optimizer.param_groups: + if group["name"] == "xyz": + p_xyz = group["params"][0] + lr_xyz = group["lr"] + break + + if p_xyz is not None: + state = self.gaussians.optimizer.state.get(p_xyz, None) + v_t = state.get("exp_avg_sq", torch.zeros_like(p_xyz)) if state else torch.zeros_like(p_xyz) + u_target = (lr_xyz / (torch.sqrt(v_t) + 1e-15)) * grad_target + u_parasitic = (lr_xyz / (torch.sqrt(v_t) + 1e-15)) * grad_parasitic + + valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gmr_trigger_status = 0.0 + if step == self.compaction_iter: + self.gaussians = subsampling(self.gaussians, self.compaction_ratio, 42, 'GMR') + self.gmr_trigger_status = 1.0 + + if step < self.opt.iterations: + if hasattr(self.gaussians, "exposure_optimizer"): + self.gaussians.exposure_optimizer.step() + self.gaussians.exposure_optimizer.zero_grad(set_to_none=True) + + if self.opt.optimizer_type == "sparse_adam": + visible = radii > 0 + self.gaussians.optimizer.step(visible, radii.shape[0]) + else: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "gmr_trigger_status": float(self.gmr_trigger_status) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + bg = torch.tensor([1, 1, 1] if self.dataset.white_background else [0, 0, 0], dtype=torch.float32, device="cuda") + render_pkg = native_render(camera, self.gaussians, self.pipe, bg, use_trained_exp=self.dataset.train_test_exp, separate_sh=(self.opt.optimizer_type == "sparse_adam")) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gigs.py b/methods/wrapper_gigs.py new file mode 100644 index 0000000000000000000000000000000000000000..e72bb88e9128311cba4324ece05ef337af0708ca --- /dev/null +++ b/methods/wrapper_gigs.py @@ -0,0 +1,448 @@ +import os +import sys +import random +import math +import torch +import numpy as np +import kornia +import nvdiffrast.torch as dr +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GI-GS_offy'))) + +from utils.loss_utils import l1_loss, ssim +from utils.image_utils import psnr +from gaussian_renderer import render +from diff_gaussian_rasterization import Gaussian_SSR +from pbr import CubemapLight, get_brdf_lut, pbr_shading +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def linear_to_srgb(linear): + eps = torch.finfo(torch.float32).eps + srgb0 = 323 / 25 * linear + srgb1 = (211 * torch.clamp(linear, min=eps) ** (5 / 12) - 11) / 200 + return torch.where(linear <= 0.0031308, srgb0, srgb1) + +def srgb_to_linear(srgb): + linear0 = 25 / 323 * srgb + linear1 = ((srgb + 0.055) / 1.055)**2.4 + return torch.where(srgb <= 0.04045, linear0, linear1) + +def get_tv_loss(gt_image, prediction, pad=1, step=1): + if pad > 1: + gt_image = F.avg_pool2d(gt_image, pad, pad) + prediction = F.avg_pool2d(prediction, pad, pad) + rgb_grad_h = torch.exp(-(gt_image[:, 1:, :] - gt_image[:, :-1, :]).abs().mean(dim=0, keepdim=True)) + rgb_grad_w = torch.exp(-(gt_image[:, :, 1:] - gt_image[:, :, :-1]).abs().mean(dim=0, keepdim=True)) + tv_h = torch.pow(prediction[:, 1:, :] - prediction[:, :-1, :], 2) + tv_w = torch.pow(prediction[:, :, 1:] - prediction[:, :, :-1], 2) + tv_loss = (tv_h * rgb_grad_h).mean() + (tv_w * rgb_grad_w).mean() + if step > 1: + for s in range(2, step + 1): + rgb_grad_h = torch.exp(-(gt_image[:, s:, :] - gt_image[:, :-s, :]).abs().mean(dim=0, keepdim=True)) + rgb_grad_w = torch.exp(-(gt_image[:, :, s:] - gt_image[:, :, :-s]).abs().mean(dim=0, keepdim=True)) + tv_h = torch.pow(prediction[:, s:, :] - prediction[:, :-s, :], 2) + tv_w = torch.pow(prediction[:, :, s:] - prediction[:, :, :-s], 2) + tv_loss += (tv_h * rgb_grad_h).mean() + (tv_w * rgb_grad_w).mean() + return tv_loss + +def get_masked_tv_loss(mask, gt_image, prediction): + rgb_grad_h = torch.exp(-(gt_image[:, 1:, :] - gt_image[:, :-1, :]).abs().mean(dim=0, keepdim=True)) + rgb_grad_w = torch.exp(-(gt_image[:, :, 1:] - gt_image[:, :, :-1]).abs().mean(dim=0, keepdim=True)) + tv_h = torch.pow(prediction[:, 1:, :] - prediction[:, :-1, :], 2) + tv_w = torch.pow(prediction[:, :, 1:] - prediction[:, :, :-1], 2) + mask = mask.float() + mask_h = mask[:, 1:, :] * mask[:, :-1, :] + mask_w = mask[:, :, 1:] * mask[:, :, :-1] + tv_loss = (tv_h * rgb_grad_h * mask_h).mean() + (tv_w * rgb_grad_w * mask_w).mean() + return tv_loss + +def get_envmap_dirs(res=[512, 1024]): + gy, gx = torch.meshgrid( + torch.linspace(0.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device="cuda"), + torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device="cuda"), + indexing="ij", + ) + sintheta, costheta = torch.sin(gy * np.pi), torch.cos(gy * np.pi) + sinphi, cosphi = torch.sin(gx * np.pi), torch.cos(gx * np.pi) + reflvec = torch.stack((sintheta * sinphi, costheta, -sintheta * cosphi), dim=-1) + return reflvec + +@register_method("gigs") +class GIGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.resolution = dataset_config.get("resolution", 1) + self.args.eval = True + self.args.metallic = True + self.args.indirect = True + self.args.pbr_iteration = 7000 + self.args.tone = False + self.args.gamma = False + self.args.normal_tv = 1.0 + self.args.brdf_tv = 1.0 + self.args.env_tv = 0.01 + self.args.radius = 0.8 + self.args.bias = 0.01 + self.args.thick = 0.05 + self.args.delta = 0.0625 + self.args.step = 16 + self.args.start = 8 + + if "360" in str(self.args.source_path).lower() and "indoor" not in str(self.args.source_path).lower(): + self.args.degree = 3 + else: + self.args.degree = 1 + + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.dataset.sh_degree = self.args.degree + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + self.brdf_lut = get_brdf_lut().cuda() + self.envmap_dirs = get_envmap_dirs() + self.cubemap = CubemapLight(base_res=256).cuda() + self.cubemap.train() + + param_groups = [{"name": "cubemap", "params": self.cubemap.parameters(), "lr": self.opt.opacity_lr}] + self.light_optimizer = torch.optim.Adam(param_groups, lr=self.opt.opacity_lr) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.canonical_rays = self.scene.get_canonical_rays() + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + try: + c2w = torch.inverse(viewpoint_cam.world_view_transform.T) + except: + return {}, {} + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + current_bg = bg if step <= self.args.pbr_iteration else torch.zeros_like(bg) + + render_pkg = render( + viewpoint_camera=viewpoint_cam, + pc=self.gaussians, + pipe=self.pipe, + bg_color=current_bg, + pad_normal=False, + derive_normal=True, + radius=self.args.radius, + bias=self.args.bias, + thick=self.args.thick, + delta=self.args.delta, + step=self.args.step, + start=self.args.start + ) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + normal_map_from_depth = render_pkg["normal_map_from_depth"] + normal_map = render_pkg["normal_map"] + albedo_map = render_pkg["albedo_map"] + roughness_map = render_pkg["roughness_map"] + metallic_map = render_pkg["metallic_map"] + + rmax, rmin = 1.0, 0.04 + roughness_map = roughness_map * (rmax - rmin) + rmin + + H, W = viewpoint_cam.image_height, viewpoint_cam.image_width + view_dirs = -((F.normalize(self.canonical_rays[:, None, :], p=2, dim=-1) * c2w[None, :3, :3]).sum(dim=-1).reshape(H, W, 3)) + + alpha_mask = viewpoint_cam.gt_alpha_mask.cuda() + gt_image = viewpoint_cam.original_image[0:3, :, :].cuda() + gt_image = (gt_image * alpha_mask + current_bg[:, None, None] * (1.0 - alpha_mask)).clamp(0.0, 1.0) + + loss_target = torch.tensor(0.0, device="cuda") + loss_parasitic = torch.tensor(0.0, device="cuda") + + if step <= self.args.pbr_iteration: + Ll1 = F.l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + ssim_val = ssim(image, gt_image) + mask = render_pkg["normal_from_depth_mask"] + normal_loss = F.l1_loss(normal_map[:, mask], normal_map_from_depth[:, mask]) + normal_tv_loss = get_tv_loss(gt_image, normal_map, pad=1, step=1) + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_val) + normal_loss + normal_tv_loss * self.args.normal_tv + loss = loss_target + loss_parasitic + else: + occlusion = render_pkg["occlusion_map"].permute(1, 2, 0) if self.args.indirect else torch.ones_like(roughness_map).permute(1, 2, 0) + normal_mask = render_pkg["normal_mask"] + out_normal_view = render_pkg["out_normal_view"] + depth_pos = render_pkg["depth_pos"] + + self.cubemap.build_mips() + pbr_result = pbr_shading( + light=self.cubemap, + normals=normal_map.permute(1, 2, 0).detach(), + view_dirs=view_dirs, + mask=normal_mask.permute(1, 2, 0), + albedo=albedo_map.permute(1, 2, 0), + roughness=roughness_map.permute(1, 2, 0), + metallic=metallic_map.permute(1, 2, 0) if self.args.metallic else None, + tone=self.args.tone, + gamma=self.args.gamma, + occlusion=occlusion.detach(), + brdf_lut=self.brdf_lut + ) + + diffuse_rgb = pbr_result["diffuse_rgb"].clamp(min=0.0, max=1.0).permute(2, 0, 1) + diffuse_rgb = torch.where(normal_mask, diffuse_rgb, current_bg[:, None, None]) + render_direct = pbr_result["render_rgb"].permute(2, 0, 1) + render_direct = torch.where(normal_mask, render_direct, current_bg[:, None, None]) + + tanfovx = math.tan(viewpoint_cam.FoVx * 0.5) + tanfovy = math.tan(viewpoint_cam.FoVy * 0.5) + SSR = Gaussian_SSR(tanfovx, tanfovy, W, H, self.args.radius, self.args.bias, self.args.thick, self.args.delta, self.args.step, self.args.start) + + if self.args.metallic: + F0 = (1.0 - metallic_map) * 0.04 + albedo_map * metallic_map + else: + F0 = torch.ones_like(albedo_map) * 0.04 + metallic_map = torch.zeros_like(roughness_map) + + linear_rgb = srgb_to_linear(render_direct) + (IRR, _) = SSR(out_normal_view.detach(), depth_pos.detach(), linear_rgb.detach(), albedo_map, roughness_map, metallic_map, F0) + IRR = linear_to_srgb(IRR) + IRR = kornia.filters.median_blur(IRR[None, ...], (3, 3))[0] + render_rgb = render_direct + IRR + + loss_target = l1_loss(render_rgb, gt_image) + + if (normal_mask == 0).sum() > 0: + brdf_tv_loss = get_masked_tv_loss(normal_mask, gt_image, torch.cat([albedo_map, roughness_map, metallic_map], dim=0)) + else: + brdf_tv_loss = get_tv_loss(gt_image, torch.cat([albedo_map, roughness_map, metallic_map], dim=0), pad=1, step=1) + + lamb_loss = (1.0 - roughness_map[normal_mask]).mean() + metallic_map[normal_mask].mean() + envmap = dr.texture(self.cubemap.base[None, ...], self.envmap_dirs[None, ...].contiguous(), filter_mode="linear", boundary_mode="cube")[0] + tv_h1 = torch.pow(envmap[1:, :, :] - envmap[:-1, :, :], 2).mean() + tv_w1 = torch.pow(envmap[:, 1:, :] - envmap[:, :-1, :], 2).mean() + env_tv_loss = tv_h1 + tv_w1 + + loss_parasitic = brdf_tv_loss * self.args.brdf_tv + lamb_loss * 0.001 + env_tv_loss * self.args.env_tv + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest, self.cubemap.base], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if step > self.args.pbr_iteration: + self.light_optimizer.zero_grad(set_to_none=True) + + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p is not None and p.grad is not None]) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if step > self.args.pbr_iteration: + self.light_optimizer.zero_grad(set_to_none=True) + + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p is not None and p.grad is not None]) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if step > self.args.pbr_iteration: + self.light_optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + if step >= self.args.pbr_iteration: + self.light_optimizer.step() + self.light_optimizer.zero_grad(set_to_none=True) + self.cubemap.clamp_(min=0.0) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "cubemap_energy": float(self.cubemap.base.norm(p=2)), + "roughness_mean": float(roughness_map.mean()), + "metallic_ratio": float((metallic_map > 0.5).float().mean()) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = render(camera, self.gaussians, self.pipe, self.background, pad_normal=False, derive_normal=True, radius=self.args.radius, bias=self.args.bias, thick=self.args.thick, delta=self.args.delta, step=self.args.step, start=self.args.start) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth_map", None)} + + def save(self, save_dir, step): + self.scene.save(step) + torch.save({"cubemap": self.cubemap.state_dict(), "light_optimizer": self.light_optimizer.state_dict(), "iteration": step}, os.path.join(save_dir, f"chkpnt{step}.pth")) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + chkpnt = os.path.join(model_path, f"chkpnt{iteration}.pth") + if os.path.exists(chkpnt): + ckpt = torch.load(chkpnt) + self.cubemap.load_state_dict(ckpt["cubemap"]) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + diff --git a/methods/wrapper_gof.py b/methods/wrapper_gof.py new file mode 100644 index 0000000000000000000000000000000000000000..5fcd793a6e24807acc84ddd6103cca8aa876afbe --- /dev/null +++ b/methods/wrapper_gof.py @@ -0,0 +1,260 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../gof_official'))) +from utils.loss_utils import l1_loss, ssim +from utils.depth_utils import depth_to_normal +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("gof") +class GOFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.trainCameras = self.scene.getTrainCameras().copy() + self.viewpoint_stack = self.trainCameras.copy() + + self.highresolution_index = [i for i, c in enumerate(self.trainCameras) if c.image_width >= 800] + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.trainCameras.copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + if random.random() < 0.3 and self.dataset.sample_more_highres and len(self.highresolution_index) > 0: + viewpoint_cam = self.trainCameras[self.highresolution_index[random.randint(0, len(self.highresolution_index) - 1)]] + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + rendering = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + image = rendering[:3, :, :] + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + distortion_map = rendering[8, :, :] if rendering.shape[0] > 8 else torch.zeros_like(image[0]) + loss_distortion = distortion_map.mean() + + depth = rendering[6, :, :] if rendering.shape[0] > 6 else torch.zeros_like(image[0]) + depth_normal, _ = depth_to_normal(viewpoint_cam, depth[None, ...]) + depth_normal = depth_normal.permute(2, 0, 1) + render_normal = rendering[3:6, :, :] if rendering.shape[0] > 5 else torch.zeros_like(image) + render_normal = F.normalize(render_normal, p=2, dim=0) + c2w = (viewpoint_cam.world_view_transform.T).inverse() + normal2 = c2w[:3, :3] @ render_normal.reshape(3, -1) + render_normal_world = normal2.reshape(3, *render_normal.shape[1:]) + normal_error = 1 - (render_normal_world * depth_normal).sum(dim=0) + loss_normal = normal_error.mean() + + lambda_dist = self.opt.lambda_distortion if step >= self.opt.distortion_from_iter else 0.0 + lambda_norm = self.opt.lambda_depth_normal if step >= self.opt.depth_normal_from_iter else 0.0 + + loss_parasitic = loss_distortion * lambda_dist + loss_normal * lambda_norm + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and loss_parasitic.item() > 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = self.gaussians.optimizer.param_groups[0]["lr"] + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + else: + u_t = grad_target + u_p = grad_parasitic + + norm_t = torch.norm(u_t, dim=1) + norm_p = torch.norm(u_p, dim=1) + valid_mask = (norm_t > 0) & (norm_p > 0) + + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(norm_p.mean() / (norm_t.mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter: + if step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "loss_distortion": float(loss_distortion), "loss_normal": float(loss_normal) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = self.gaussians.get_opacity_with_3D_filter.clone().detach() + scales = self.gaussians.get_scaling_with_3D_filter.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + return {"image": render_pkg["render"][:3, :, :], "depth": render_pkg["render"][6, :, :] if render_pkg["render"].shape[0] > 6 else None} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling_with_3D_filter + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.gaussians.get_opacity_with_3D_filter)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = self.gaussians.get_opacity_with_3D_filter.squeeze() + scales = self.gaussians.get_scaling_with_3D_filter + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + diff --git a/methods/wrapper_gsdf.py b/methods/wrapper_gsdf.py new file mode 100644 index 0000000000000000000000000000000000000000..c72f0db7a565646c3ec74ec4eb22c52d68d30a65 --- /dev/null +++ b/methods/wrapper_gsdf.py @@ -0,0 +1,261 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +# 强制路径隔离：将 GSDF 路径插入最前端并清理冲突模块 +base_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GSDF_run')) +gs_path = os.path.join(base_path, "gaussian_splatting") +sys.path.insert(0, base_path) +sys.path.insert(0, gs_path) + +for mod in ["gaussian_renderer", "scene", "arguments", "utils", "gaussian_renderer.render"]: + if mod in sys.modules: + del sys.modules[mod] + +from core.registry import register_method +from core.base_method import BaseMethod + +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from gaussian_renderer import prefilter_voxel, generate_neural_gaussians +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("gsdf") +class GSDFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.feat_dim, self.dataset.n_offsets, self.dataset.voxel_size, self.dataset.update_depth, self.dataset.update_init_factor, self.dataset.update_hierachy_factor, self.dataset.use_feat_bank, self.dataset.use_tcnn) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_anchors = self.gaussians.get_anchor.shape[0] + + self.cached_xyz = None + self.cached_scaling = None + self.cached_opacity = None + self.cached_rot = None + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + voxel_visible_mask = prefilter_voxel(viewpoint_cam, self.gaussians, self.pipe, self.background) + retain_grad = (step < 15000 and step >= 0) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask, retain_grad=retain_grad) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + if retain_grad and viewspace_point_tensor is not None: + viewspace_point_tensor.retain_grad() + visibility_filter = render_pkg["visibility_filter"] + offset_selection_mask = render_pkg["selection_mask"] + scaling = render_pkg["scaling"] + opacity = render_pkg["neural_opacity"] + + self.cached_xyz = render_pkg.get("xyz", None) + self.cached_scaling = scaling + self.cached_opacity = opacity + self.cached_rot = render_pkg.get("rotations", None) + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + loss_parasitic = 0.01 * scaling.prod(dim=1).mean() + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._scaling.grad.clone() if self.gaussians._scaling.grad is not None else torch.zeros_like(self.gaussians._scaling) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._scaling.grad.clone() if self.gaussians._scaling.grad is not None else torch.zeros_like(self.gaussians._scaling) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.update_until and step > self.opt.start_stat: + self.gaussians.training_statis(viewspace_point_tensor, opacity, visibility_filter, offset_selection_mask, voxel_visible_mask) + if step > self.opt.update_from and step % self.opt.update_interval == 0: + self.gaussians.adjust_anchor(extent=self.scene.cameras_extent, check_interval=self.opt.update_interval, success_threshold=self.opt.success_threshold, grad_threshold=self.opt.densify_grad_threshold, min_opacity=self.opt.min_opacity) + elif step == self.opt.update_until: + if hasattr(self.gaussians, 'opacity_accum'): del self.gaussians.opacity_accum + if hasattr(self.gaussians, 'offset_gradient_accum'): del self.gaussians.offset_gradient_accum + if hasattr(self.gaussians, 'offset_denom'): del self.gaussians.offset_denom + torch.cuda.empty_cache() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_anchors = self.gaussians.get_anchor.shape[0] + num_gaussians = scaling.shape[0] if scaling is not None else 0 + + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "num_anchors": int(num_anchors), + "delta_N": int(num_anchors - self.last_n_anchors), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "anchor_occupancy_ratio": float(num_gaussians / (num_anchors * self.dataset.n_offsets + 1e-7)), + "loss_scaling_reg": float(loss_parasitic) + } + self.last_n_anchors = num_anchors + + histograms = {} + if step % 1000 == 0 and self.cached_scaling is not None: + histograms["opacity"] = self.cached_opacity.clone().detach() + histograms["scaling"] = self.cached_scaling.clone().detach() + scales_2d = self.cached_scaling[:, :2] if self.cached_scaling.shape[1] >= 2 else self.cached_scaling + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + voxel_visible_mask = prefilter_voxel(camera, self.gaussians, self.pipe, self.background) + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth_hand", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply_sparse_gaussian(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + self.gaussians.load_mlp_checkpoints(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'checkpoint.pth')) + + def get_spatial_centers(self): + return self.gaussians.get_anchor + + def compute_physical_metrics(self, cameras=None): + metrics = {} + if self.cached_scaling is None or self.cached_opacity is None: + return metrics + + with torch.no_grad(): + scales = self.cached_scaling + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.cached_opacity)) + + if cameras is not None and len(cameras) > 0 and self.cached_rot is not None: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.cached_rot.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + + if cameras is None or len(cameras) == 0: + cam = self.scene.getTrainCameras()[0] + else: + cam = cameras[0] + + xyz, _, opacity, scaling, _, _, _ = generate_neural_gaussians(cam, self.gaussians, None, is_training=True) + opacity = opacity.squeeze() + + sigma_sq = (scaling[:, :2].max(dim=1)[0].pow(2)) if scaling.shape[1] >= 2 else scaling.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + if N_gaussians == 0: + return densities + + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacity.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gsgn.py b/methods/wrapper_gsgn.py new file mode 100644 index 0000000000000000000000000000000000000000..9527bb64931ba293225a3030c320c7554732bdf7 --- /dev/null +++ b/methods/wrapper_gsgn.py @@ -0,0 +1,328 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/3DGS-LM_offy') +from utils.loss_utils import l1_loss, ssim, l2_loss +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams, GSGNParams +from train import lm_step + +@register_method("gsgn") +class GSGNWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.gn = GSGNParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.gsgn = self.gn.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.viewpoint_stack_indices = list(range(len(self.viewpoint_stack))) + if self.gsgn.image_subsample_frame_selection_mode == "random": + random.shuffle(self.viewpoint_stack_indices) + + self.viewpoint_stack_sgd = self.scene.getTrainCameras().copy() + + self.trust_region_radius = self.gsgn.trust_region_radius + self.radius_decrease_factor = self.gsgn.radius_decrease_factor + self.forward_output = None + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _sgd_step(self, step): + self.gaussians.update_learning_rate(step) + if self.opt.increase_SH_deg_every > 0 and (step % self.opt.increase_SH_deg_every) == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack_sgd: + self.viewpoint_stack_sgd = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack_sgd.pop(random.randint(0, len(self.viewpoint_stack_sgd)-1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + loss = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + if self.opt.lambda_l2 > 0: + loss = loss + self.opt.lambda_l2 * l2_loss(image, gt_image) + + loss.backward() + + with torch.no_grad(): + if step > -1 and step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + return float(loss), float(Ll1), float(1.0 - ssim(image, gt_image)) + + def train_iteration(self, step): + metrics = {} + histograms = {} + + if self.gsgn.num_sgd_iterations_before_gn > 0 and step == 1: + for i in range(self.gsgn.num_sgd_iterations_before_gn): + self._sgd_step(1 + i) + + timing_dict = {} + out_dict = lm_step( + gaussians=self.gaussians, + viewpoint_stack=self.viewpoint_stack, + viewpoint_stack_indices=self.viewpoint_stack_indices, + opt=self.opt, + pipe=self.pipe, + background=self.background, + gsgn=self.gsgn, + iteration=step, + trust_region_radius=self.trust_region_radius, + radius_decrease_factor=self.radius_decrease_factor, + timing_dict=timing_dict, + forward_output=self.forward_output + ) + + self.trust_region_radius = out_dict["trust_region_radius"] + self.radius_decrease_factor = out_dict["radius_decrease_factor"] + self.forward_output = out_dict["forward_output"] + + loss_v = float(out_dict.get("loss", 0.0)) + loss_l1_v = 0.0 + loss_ssim_v = 0.0 + + if self.gsgn.num_sgd_iterations_between_gn > 0 and (step % self.gsgn.sgd_between_gn_every) == 0: + for i in range(self.gsgn.num_sgd_iterations_between_gn): + l_total, l1_val, ssim_val = self._sgd_step(-1) + loss_v = l_total + loss_l1_v = l1_val + loss_ssim_v = ssim_val + + if self.gsgn.num_sgd_iterations_after_gn > 0 and step == self.opt.iterations: + for i in range(self.gsgn.num_sgd_iterations_after_gn): + self._sgd_step(-1) + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + viewpoint_cam = self.viewpoint_stack[0] + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + gt_image = viewpoint_cam.original_image.cuda() + + loss_target = (1.0 - self.opt.lambda_dssim) * l1_loss(image, gt_image) + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get('exp_avg_sq', torch.ones_like(p)) + eff_step = p.grad.clone() / (torch.sqrt(v) + 1e-8) + group_grads.append(eff_step.reshape(-1)) + grads_target[group_name] = torch.cat(group_grads) if group_grads else torch.empty(0, device="cuda") + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get('exp_avg_sq', torch.ones_like(p)) + eff_step = p.grad.clone() / (torch.sqrt(v) + 1e-8) + group_grads.append(eff_step.reshape(-1)) + grads_parasitic[group_name] = torch.cat(group_grads) if group_grads else torch.empty(0, device="cuda") + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + metrics[f"sti_{group_name}"] = ti + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics.update({ + "loss": loss_v, + "loss_l1": loss_l1_v, + "loss_ssim": loss_ssim_v, + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + }) + + log_info = out_dict.get("log_info", {}) + pcg_info = log_info.get("pcg_info", {}) + if pcg_info: + metrics["pcg_niter"] = float(pcg_info.get("niter", 0)) + metrics["pcg_err"] = float(pcg_info.get("err", 0)) + metrics["lm_cost_change"] = float(log_info.get("cost_change", 0.0)) + metrics["trust_region_radius"] = float(self.trust_region_radius) + + self.last_n_gaussians = num_gaussians + + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gslpm.py b/methods/wrapper_gslpm.py new file mode 100644 index 0000000000000000000000000000000000000000..f6a37db46af0e21d6948b1be50d2528d4be97938 --- /dev/null +++ b/methods/wrapper_gslpm.py @@ -0,0 +1,246 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GS-LPM'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from lpm.lpm import LPM + +@register_method("gslpm") +class GSLPMWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + self.lpm = LPM(self.scene, self.gaussians, angle=self.opt.angle) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_effective_step(self, param, grad, group_name): + lr = 0.0 + if self.gaussians.optimizer is not None: + for group in self.gaussians.optimizer.param_groups: + if group["name"] == group_name: + lr = group["lr"] + break + state = self.gaussians.optimizer.state.get(param, None) + if state is not None and "exp_avg_sq" in state: + vt = state["exp_avg_sq"] + return (lr / (torch.sqrt(vt) + 1e-15)) * grad + return lr * grad + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + current_view_index, sampled_index = self.lpm.find_neighbor_cam(viewpoint_cam) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target_spatial = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic_spatial = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + u_target = self._get_effective_step(self.gaussians._xyz, grad_target_spatial, "xyz") + u_parasitic = self._get_effective_step(self.gaussians._xyz, grad_parasitic_spatial, "xyz") + + valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + lpm_addition_triggered = 0 + lpm_calibration_triggered = 0 + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.lpm.points_addition(self.opt.densify_grad_threshold, size_threshold, viewpoint_cam, current_view_index, sampled_index, image, gt_image) + lpm_addition_triggered = 1 + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step > self.opt.reset_from_iter and step % self.opt.reset_interval == 0 and step < self.opt.reset_until_iter: + self.lpm.points_calibration(self.opt.densify_grad_threshold, viewpoint_cam, current_view_index, sampled_index, image, gt_image) + lpm_calibration_triggered = 1 + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "lpm_addition_triggered": int(lpm_addition_triggered), "lpm_calibration_triggered": int(lpm_calibration_triggered) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.gaussians._scaling) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_gspull.py b/methods/wrapper_gspull.py new file mode 100644 index 0000000000000000000000000000000000000000..80aafd965f3af372b84e4298096c76008d18d813 --- /dev/null +++ b/methods/wrapper_gspull.py @@ -0,0 +1,223 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +GS_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GS-Pull')) +# 强隔离：清空残留路径并确保 GS-Pull 绝对优先 +if GS_ROOT not in sys.path: + sys.path.insert(0, GS_ROOT) +# GS-Pull's gaussian_splatting is a vanilla 3DGS copy — its scene/__init__.py +# does `from utils.system_utils import ...` expecting the gaussian_splatting/ +# subdir to be on sys.path. +_GS_INNER = os.path.join(GS_ROOT, 'gaussian_splatting') +if _GS_INNER not in sys.path: + sys.path.insert(0, _GS_INNER) + +from gaussian_splatting.scene import Scene, GaussianModel +from gaussian_splatting.arguments import ModelParams, PipelineParams, OptimizationParams as GSOptParams +from gaussian_splatting.gaussian_renderer import render as vanilla_render + +from sugar_scene.sugar_model import SuGaR +from sugar_scene.sugar_optimizer import OptimizationParams as SuGaROptParams, SuGaROptimizer +from sugar_scene.sugar_densifier import SuGaRDensifier +from sugar_utils.loss_utils import ssim, l1_loss +from np_utils.train import Runner + +@register_method("gspull") +class GSPullWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.source_path = dataset_config["source_path"] + self.model_path = dataset_config["model_path"] + self.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + parser = ArgumentParser() + lp = ModelParams(parser) + op = GSOptParams(parser) + pp = PipelineParams(parser) + args = parser.parse_args([]) + args.source_path = self.source_path + args.model_path = self.model_path + args.resolution = self.resolution + args.eval = True + + self.dataset = lp.extract(args) + self.gs_opt = op.extract(args) + self.pipe = pp.extract(args) + + # 防呆逻辑：兼容某些变体改名为 shfeature_lr 的情况 + if not hasattr(self.gs_opt, 'shfeature_lr') and hasattr(self.gs_opt, 'feature_lr'): + setattr(self.gs_opt, 'shfeature_lr', getattr(self.gs_opt, 'feature_lr')) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.gs_opt) + + self.bg_color = torch.tensor([1, 1, 1] if self.dataset.white_background else [0, 0, 0], dtype=torch.float32, device="cuda") + self.train_cameras = self.scene.getTrainCameras().copy() + self.viewpoint_stack = self.train_cameras.copy() + + self.sugar = None + self.optimizer = None + self.densifier = None + self._iter = 0 + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _init_sugar_stage(self): + self.sugar = SuGaR( + nerfmodel=None, + points=self.gaussians.get_xyz.detach(), + colors=torch.zeros_like(self.gaussians.get_xyz), + initialize=False, + sh_levels=self.gaussians.active_sh_degree + 1, + learnable_positions=True, + triangle_scale=1.0, + keep_track_of_knn=True, + knn_to_track=16, + beta_mode="average", + freeze_gaussians=False, + surface_mesh_to_bind=None, + ) + with torch.no_grad(): + self.sugar._points.copy_(self.gaussians._xyz) + self.sugar._scales.copy_(self.gaussians._scaling) + self.sugar._quaternions.copy_(self.gaussians._rotation) + self.sugar.all_densities.copy_(self.gaussians._opacity) + self.sugar._sh_coordinates_dc.copy_(self.gaussians._features_dc) + self.sugar._sh_coordinates_rest.copy_(self.gaussians._features_rest) + + self.sugar.part_num = 1 + self.sugar.neus = Runner(self.model_path, None, part_num=self.sugar.part_num) + + spatial_lr_scale = self.scene.cameras_extent + opt_p = SuGaROptParams( + iterations=15000, position_lr_init=0.00016, position_lr_final=0.0000016, + position_lr_max_steps=30000, scaling_lr=0.005, opacity_lr=0.05, + feature_lr=0.0025, rotation_lr=0.001 + ) + self.optimizer = SuGaROptimizer(self.sugar, opt_p, spatial_lr_scale=spatial_lr_scale) + self.densifier = SuGaRDensifier(self.sugar, self.optimizer, 0.0002, 0.005, 20, spatial_lr_scale, 0.01) + + def train_iteration(self, step): + self._iter += 1 + metrics = {} + histograms = {} + + if self._iter <= 7000: + if not self.viewpoint_stack: + self.viewpoint_stack = self.train_cameras.copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = vanilla_render(viewpoint_cam, self.gaussians, self.pipe, self.bg_color) + image, viewspace_point_tensor, visibility_filter, radii = render_pkg["render"], render_pkg["viewspace_points"], render_pkg["visibility_filter"], render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + loss = 0.8 * l1_loss(image, gt_image) + 0.2 * (1.0 - ssim(image, gt_image)) + loss.backward() + + with torch.no_grad(): + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if self._iter > 500 and self._iter % 100 == 0: + self.gaussians.densify_and_prune(self.gs_opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, 20) + if self._iter % 3000 == 0: + self.gaussians.reset_opacity() + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = len(self.gaussians.get_xyz) + metrics = { + "loss": float(loss), "num_gaussians": int(num_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) + } + self.last_n_gaussians = num_gaussians + return metrics, histograms + + if self._iter == 7001: + self._init_sugar_stage() + + is_pulling = self._iter > 9000 + if self._iter == 9001: + prune_mask = (self.sugar.strengths < 0.5).squeeze() + self.densifier.prune_points(prune_mask) + self.sugar.reset_neighbors() + self.sugar.neus.reset_datasets(self.model_path, self.sugar.points.detach().cpu().numpy(), iteration=9000, scene_name="scene") + + if not self.viewpoint_stack: + self.viewpoint_stack = self.train_cameras.copy() + cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + outputs = self.sugar.render_image_gaussian_rasterizer( + camera_indices=self.train_cameras.index(cam), bg_color=self.bg_color, + sh_deg=self.sugar.sh_levels-1, compute_covariance_in_rasterizer=True, return_2d_radii=True + ) + + pred_rgb = outputs['image'].view(-1, cam.image_height, cam.image_width, 3).transpose(-1, -2).transpose(-2, -3) + gt_rgb = cam.original_image.cuda().unsqueeze(0) + loss = 0.8 * l1_loss(pred_rgb, gt_rgb) + 0.2 * (1.0 - ssim(pred_rgb, gt_rgb)) + loss.backward() + + self.optimizer.step() + self.optimizer.zero_grad(set_to_none=True) + + metrics = {"loss": float(loss), "num_gaussians": len(self.sugar.points)} + return metrics, {} + + def render(self, camera): + with torch.no_grad(): + if self._iter <= 7000: + return {"image": vanilla_render(camera, self.gaussians, self.pipe, self.bg_color)["render"], "depth": None} + idx = 0 + for i, c in enumerate(self.train_cameras): + if c.uid == camera.uid: idx = i; break + outputs = self.sugar.render_image_gaussian_rasterizer(camera_indices=idx) + return {"image": outputs["image"], "depth": None} + + def save(self, save_dir, step): + if self.sugar: + self.sugar.save_model(path=os.path.join(save_dir, f"{step}.pt"), iteration=step) + else: + self.gaussians.save_ply(os.path.join(save_dir, "point_cloud", f"iteration_{step}", "point_cloud.ply")) + + def load(self, model_path, iteration): + pass + + def get_spatial_centers(self): + return self.sugar.points if self.sugar else self.gaussians.get_xyz + + def compute_physical_metrics(self, cameras=None): + return {"num_gaussians": float(len(self.get_spatial_centers()))} + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + return torch.zeros(len(query_points), device="cuda") diff --git a/methods/wrapper_gsshader.py b/methods/wrapper_gsshader.py new file mode 100644 index 0000000000000000000000000000000000000000..a1725a3dd6aa4ced3456217fa0aa11abd4fa9e63 --- /dev/null +++ b/methods/wrapper_gsshader.py @@ -0,0 +1,256 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +# [Isolation Step 1] 剔除所有路径污染 +sys.path = [p for p in sys.path if 'taming-3dgs' not in p] + +# [Isolation Step 2] 插入正确路径 +gsshader_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../gsshader_official')) +if gsshader_root not in sys.path: + sys.path.insert(0, gsshader_root) + +# [Isolation Step 3] 强制清理模块缓存，防止重名污染 +for m in list(sys.modules.keys()): + base_m = m.split('.')[0] + if base_m in ["utils", "scene", "gaussian_renderer", "arguments"]: + del sys.modules[m] + +from utils.loss_utils import l1_loss, ssim, predicted_normal_loss, delta_normal_loss, zero_one_loss +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("gsshader") +class GSShaderWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args, _ = self.parser.parse_known_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel( + self.dataset.sh_degree, + self.dataset.brdf_dim, + self.dataset.brdf_mode, + self.dataset.brdf_envmap_res + ) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + if self.pipe.brdf: + self.gaussians.set_requires_grad("normal", state=step >= self.opt.normal_reg_from_iter) + self.gaussians.set_requires_grad("normal2", state=step >= self.opt.normal_reg_from_iter) + if self.gaussians.brdf_mode == "envmap": + self.gaussians.brdf_mlp.build_mips() + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, debug=False) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + loss_parasitic = torch.tensor(0.0, device="cuda") + if self.pipe.brdf and step > self.opt.normal_reg_from_iter: + if step < self.opt.normal_reg_util_iter: + loss_parasitic += self.opt.lambda_predicted_normal * predicted_normal_loss(render_pkg["normal"], render_pkg["normal_ref"], render_pkg["alpha"]) + loss_parasitic += self.opt.lambda_zero_one * zero_one_loss(render_pkg["alpha"]) + if "delta_normal_norm" in render_pkg.keys(): + loss_parasitic += self.opt.lambda_delta_reg * delta_normal_loss(render_pkg["delta_normal_norm"], render_pkg["alpha"]) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and loss_parasitic.item() > 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + if self.pipe.brdf and self.pipe.brdf_mode == "envmap": + self.gaussians.brdf_mlp.clamp_(min=0.0, max=1.0) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, debug=False) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if self.pipe.brdf: + metrics["roughness_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._roughness + self.gaussians.roughness_bias))) + metrics["specular_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._specular))) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_hac_plus.py b/methods/wrapper_hac_plus.py new file mode 100644 index 0000000000000000000000000000000000000000..f850d61f5924cec31851b7764c75df5c00ea1f7d --- /dev/null +++ b/methods/wrapper_hac_plus.py @@ -0,0 +1,251 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../hac_plus_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import prefilter_voxel, render as native_render, generate_neural_gaussians +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from utils.encodings import get_binary_vxl_size + +@register_method("hac_plus") +class HACPlusWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.args.n_features = 4 + self.args.log2 = 13 + self.args.log2_2D = 15 + self.args.lmbda = 0.001 + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + is_synthetic_nerf = os.path.exists(os.path.join(self.args.source_path, "transforms_train.json")) + + self.gaussians = GaussianModel( + self.dataset.feat_dim, + self.dataset.n_offsets, + self.dataset.voxel_size, + self.dataset.update_depth, + self.dataset.update_init_factor, + self.dataset.update_hierachy_factor, + self.dataset.use_feat_bank, + n_features_per_level=self.args.n_features, + log2_hashmap_size=self.args.log2, + log2_hashmap_size_2D=self.args.log2_2D, + is_synthetic_nerf=is_synthetic_nerf, + ) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.update_anchor_bound() + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = self.gaussians.get_anchor.shape[0] * self.gaussians.n_offsets + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + voxel_visible_mask = prefilter_voxel(viewpoint_cam, self.gaussians, self.pipe, self.background) + retain_grad = (step < self.opt.update_until and step >= 0) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask, retain_grad=retain_grad, step=step) + image = render_pkg["render"] + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + scaling_reg = render_pkg["scaling"].prod(dim=1).mean() + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + 0.01 * scaling_reg + + bit_per_param = render_pkg.get("bit_per_param", None) + bit_hash_grid_val = 0.0 + if bit_per_param is not None: + _, bit_hash_grid, _, _ = get_binary_vxl_size((self.gaussians.get_encoding_params() + 1) / 2) + bit_hash_grid_val = float(bit_hash_grid) + denom = self.gaussians._anchor.shape[0] * (self.gaussians.feat_dim + 6 + 3 * self.gaussians.n_offsets) + loss_parasitic = loss_parasitic + self.args.lmbda * (bit_per_param + bit_hash_grid / denom) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_t = self.gaussians._anchor.grad.clone() if self.gaussians._anchor.grad is not None else torch.zeros_like(self.gaussians._anchor) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_p = self.gaussians._anchor.grad.clone() if self.gaussians._anchor.grad is not None else torch.zeros_like(self.gaussians._anchor) + + state = self.gaussians.optimizer.state.get(self.gaussians._anchor, {}) + v = state.get("exp_avg_sq", torch.ones_like(grad_t) * 1e-8) + lr = 0.0 + for pg in self.gaussians.optimizer.param_groups: + if pg["name"] == "anchor": + lr = pg["lr"] + break + + u_t = (lr / (torch.sqrt(v) + 1e-8)) * grad_t + u_p = (lr / (torch.sqrt(v) + 1e-8)) * grad_p + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.update_until and step > self.opt.start_stat: + self.gaussians.training_statis(render_pkg["viewspace_points"], render_pkg["neural_opacity"], render_pkg["visibility_filter"], render_pkg["selection_mask"], voxel_visible_mask) + if step not in range(3000, 4000): + if step > self.opt.update_from and step % self.opt.update_interval == 0: + self.gaussians.adjust_anchor(check_interval=self.opt.update_interval, success_threshold=self.opt.success_threshold, grad_threshold=self.opt.densify_grad_threshold, min_opacity=self.opt.min_opacity) + elif step == self.opt.update_until: + if hasattr(self.gaussians, "opacity_accum"): + del self.gaussians.opacity_accum + if hasattr(self.gaussians, "offset_gradient_accum"): + del self.gaussians.offset_gradient_accum + if hasattr(self.gaussians, "offset_denom"): + del self.gaussians.offset_denom + torch.cuda.empty_cache() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_anchor.shape[0] * self.gaussians.n_offsets + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "bit_per_param": float(bit_per_param) if bit_per_param is not None else 0.0, + "bit_hash_grid": float(bit_hash_grid_val) + } + self.last_n_gaussians = num_gaussians + + return metrics, {} + + def render(self, camera): + with torch.no_grad(): + voxel_visible_mask = prefilter_voxel(camera, self.gaussians, self.pipe, self.background) + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.scene = Scene(self.dataset, self.gaussians, load_iteration=iteration, shuffle=False) + self.gaussians.eval() + + def get_spatial_centers(self): + return self.gaussians.get_anchor + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + cam = cameras[0] if cameras and len(cameras) > 0 else self.viewpoint_stack[0] + xyz, color, opacity, scaling, rot, _ = generate_neural_gaussians(cam, self.gaussians, visible_mask=None, is_training=False, step=30000) + + scales_2d = scaling[:, :2] if scaling.dim() > 1 and scaling.shape[1] >= 2 else scaling.unsqueeze(-1).expand(-1, 2) + max_s, _ = torch.max(scales_2d, dim=1) + min_s, _ = torch.min(scales_2d, dim=1) + gamma = max_s / (min_s + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(opacity)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(rot.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + v = query_points.shape[0] + densities = torch.zeros(v, device="cuda") + cam = cameras[0] if cameras and len(cameras) > 0 else self.viewpoint_stack[0] + xyz, color, opacity, scaling, rot, _ = generate_neural_gaussians(cam, self.gaussians, visible_mask=None, is_training=False, step=30000) + + opacities = opacity.squeeze() + sigma_sq = (scaling[:, :2].max(dim=1)[0].pow(2)) if scaling.shape[1] >= 2 else scaling.squeeze().pow(2) + + n_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (n_gaussians + 1)) + for i in range(0, v, chunk_size): + end = min(i + chunk_size, v) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_hogs.py b/methods/wrapper_hogs.py new file mode 100644 index 0000000000000000000000000000000000000000..a0b621cced0caf64f5e263071b8cadc4a54acbfb --- /dev/null +++ b/methods/wrapper_hogs.py @@ -0,0 +1,252 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../HoGS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("hogs") +class HoGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.iterations = 50000 + self.opt.densify_until_iter = 30000 + self.opt.opacity_reset_interval = 6000 + self.opt.w_lr = 0.0002 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + loss_target = (1.0 - self.opt.lambda_dssim) * l1_loss(image, gt_image) + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + def get_eff_step(l_val): + self.gaussians.optimizer.zero_grad(set_to_none=True) + l_val.backward(retain_graph=True) + steps = [] + for name in ["xyz", "w"]: + for group in self.gaussians.optimizer.param_groups: + if group["name"] == name: + p = group["params"][0] + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + s = (group["lr"] / (torch.sqrt(v) + 1e-15)) * p.grad.clone() + else: + s = group["lr"] * p.grad.clone() + steps.append(s.view(p.shape[0], -1)) + return torch.cat(steps, dim=1) if len(steps) > 0 else torch.zeros(self.gaussians.get_xyz.shape[0], 4, device="cuda") + + step_target = get_eff_step(loss_target) + step_parasitic = get_eff_step(loss_parasitic) + + norm_t = torch.norm(step_target, dim=1) + norm_p = torch.norm(step_parasitic, dim=1) + valid_mask = (norm_t > 0) & (norm_p > 0) + + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(step_target[valid_mask], step_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(norm_p.mean() / (norm_t.mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + w_inv = (1.0 / torch.exp(self.gaussians._w)).clone().detach() + raw_scales = torch.exp(self.gaussians._scaling).clone().detach() + eff_scales = raw_scales * w_inv.unsqueeze(-1) + histograms["scaling"] = eff_scales + scales_2d = eff_scales[:, :2] if eff_scales.shape[1] >= 2 else eff_scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + w_inv = 1.0 / torch.exp(self.gaussians._w) + return self.gaussians._xyz * w_inv.unsqueeze(-1) + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + w_inv = 1.0 / torch.exp(self.gaussians._w) + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) * w_inv.unsqueeze(-1) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + w_val = torch.exp(self.gaussians._w) + metrics["w_mean"] = float(torch.mean(w_val)) + metrics["w_median"] = float(torch.median(w_val)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + + w_inv = 1.0 / torch.exp(self.gaussians._w) + xyz = self.gaussians._xyz * w_inv.unsqueeze(-1) + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) * w_inv.unsqueeze(-1) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_improvingadc.py b/methods/wrapper_improvingadc.py new file mode 100644 index 0000000000000000000000000000000000000000..83d93b38517c1800753d5234cbe404258374f9c4 --- /dev/null +++ b/methods/wrapper_improvingadc.py @@ -0,0 +1,303 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/Improving-ADC-3DGS_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("improvingadc") +class ImprovingADCWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + self.gaussians.set_dl(folder=self.args.model_path, log_frq=50, param_log=False) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + self.dense_iters = ((self.opt.densify_until_iter - self.opt.densify_from_iter) // self.opt.densification_interval) + self.start_thresh = 0.0001 + self.end_thresh = 0.0004 + self.dense_factor = pow((self.end_thresh / self.start_thresh), pow(self.dense_iters, -1)) + self.current_dense_thresh = self.start_thresh + self.per_gaussian_alpha = None + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + self.gaussians.dl.iteration = step + self.gaussians.dl.log_in_file(self.gaussians) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, depth_threshold=self.opt.depth_threshold * self.scene.old_extent) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + fake_color = torch.zeros_like(self.gaussians._xyz, requires_grad=True) + fake_render = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, override_color=fake_color, depth_threshold=self.opt.depth_threshold * self.scene.old_extent)["render"] + fake_loss = torch.sum(fake_render.view(-1)) + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + fake_loss + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target_xyz = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + grads_target = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get("exp_avg_sq", torch.zeros_like(p.grad)) + for pg in self.gaussians.optimizer.param_groups: + if pg['params'][0] is p: + lr = pg['lr'] + break + else: + lr = 1e-4 + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + group_grads.append(u.reshape(-1)) + if group_grads: + grads_target[group_name] = torch.cat(group_grads) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic_xyz = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get("exp_avg_sq", torch.zeros_like(p.grad)) + for pg in self.gaussians.optimizer.param_groups: + if pg['params'][0] is p: + lr = pg['lr'] + break + else: + lr = 1e-4 + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + group_grads.append(u.reshape(-1)) + if group_grads: + grads_parasitic[group_name] = torch.cat(group_grads) + + valid_mask = (torch.norm(grad_target_xyz, dim=1) > 0) & (torch.norm(grad_parasitic_xyz, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target_xyz[valid_mask], grad_parasitic_xyz[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic_xyz, dim=1).mean() / (torch.norm(grad_target_xyz, dim=1).mean() + 1e-7)) + + for group_name in param_groups_map: + gt = grads_target.get(group_name) + gp = grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + if self.per_gaussian_alpha is not None: + self.per_gaussian_alpha += torch.mean(fake_color.grad, dim=1) + else: + self.per_gaussian_alpha = torch.mean(fake_color.grad, dim=1) + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter, render_pkg["pixels"]) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.current_dense_thresh *= self.dense_factor + self.gaussians.densify_and_prune(self.current_dense_thresh, 0.005, self.scene.cameras_extent, size_threshold, metric=self.per_gaussian_alpha) + self.per_gaussian_alpha = None + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + alpha_mean_val = float(self.per_gaussian_alpha.mean()) if self.per_gaussian_alpha is not None else 0.0 + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "dense_threshold_current": float(self.current_dense_thresh), + "per_gaussian_alpha_mean": alpha_mean_val + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + bg = torch.tensor([1, 1, 1] if self.dataset.white_background else [0, 0, 0], dtype=torch.float32, device="cuda") + render_pkg = native_render(camera, self.gaussians, self.pipe, bg) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_lapisgs.py b/methods/wrapper_lapisgs.py new file mode 100644 index 0000000000000000000000000000000000000000..8abaffb6690bbe03e1e805f9ecf351af070a182e --- /dev/null +++ b/methods/wrapper_lapisgs.py @@ -0,0 +1,270 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../LapisGS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("lapisgs") +class LapisGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dynamic_opacity = hyperparams.get("dynamic_opacity", True) + self.foundation_gs_path = hyperparams.get("foundation_gs_path", None) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.lambda_dssim = 0.8 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + + if self.foundation_gs_path and os.path.exists(self.foundation_gs_path): + self.gaussians.active_sh_degree = self.gaussians.max_sh_degree + self.foundation_gs = GaussianModel(self.dataset.sh_degree) + self.foundation_gs.load_foundation_gs(self.foundation_gs_path) + + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + + if self.foundation_gs_path and os.path.exists(self.foundation_gs_path): + self.gaussians = self._gs_merge(self.foundation_gs, self.gaussians) + self.gaussians.training_setup(self.opt) + self.gaussians.size_fixedGS = self.foundation_gs._xyz.shape[0] + if self.dynamic_opacity: + self.gaussians.if_dynamic_opacity = True + self.gaussians.trainable_attributes = ["opacity"] + else: + self.gaussians.trainable_attributes = [] + else: + self.gaussians.training_setup(self.opt) + self.gaussians.size_fixedGS = None + self.gaussians.if_dynamic_opacity = False + self.gaussians.trainable_attributes = [] + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _gs_merge(self, foundation_gs, train_gs): + merged_gaussians = GaussianModel(foundation_gs.max_sh_degree) + merged_gaussians._xyz = torch.nn.Parameter(torch.cat((foundation_gs._xyz, train_gs._xyz), dim=0)) + merged_gaussians._features_dc = torch.nn.Parameter(torch.cat((foundation_gs._features_dc, train_gs._features_dc), dim=0)) + merged_gaussians._features_rest = torch.nn.Parameter(torch.cat((foundation_gs._features_rest, train_gs._features_rest), dim=0)) + merged_gaussians._scaling = torch.nn.Parameter(torch.cat((foundation_gs._scaling, train_gs._scaling), dim=0)) + merged_gaussians._rotation = torch.nn.Parameter(torch.cat((foundation_gs._rotation, train_gs._rotation), dim=0)) + merged_gaussians._opacity = torch.nn.Parameter(torch.cat((foundation_gs._opacity, train_gs._opacity), dim=0)) + merged_gaussians.max_radii2D = torch.cat((foundation_gs.max_radii2D.clone().detach(), train_gs.max_radii2D.clone().detach()), dim=0) + return merged_gaussians + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._opacity.grad.clone() if self.gaussians._opacity.grad is not None else torch.zeros_like(self.gaussians._opacity) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._opacity.grad.clone() if self.gaussians._opacity.grad is not None else torch.zeros_like(self.gaussians._opacity) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0 and self.gaussians.get_xyz.shape[0] < 6_000_000: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if self.gaussians.size_fixedGS: + self.gaussians.freeze_fixedGS() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + base_alpha = 0.0 + enh_alpha = 0.0 + fixed_ratio = 0.0 + if self.gaussians.size_fixedGS is not None: + fixed_gs = self.gaussians.size_fixedGS + base_alpha = float(torch.sigmoid(self.gaussians._opacity[:fixed_gs]).mean()) + if num_gaussians > fixed_gs: + enh_alpha = float(torch.sigmoid(self.gaussians._opacity[fixed_gs:]).mean()) + fixed_ratio = float(fixed_gs / num_gaussians) + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "base_layer_alpha_mean": base_alpha, + "enhancement_layer_alpha_mean": enh_alpha, + "fixed_gs_ratio": fixed_ratio + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_lightgaussian.py b/methods/wrapper_lightgaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..8ba9c150929399715b02007e514041185120afb4 --- /dev/null +++ b/methods/wrapper_lightgaussian.py @@ -0,0 +1,268 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../LightGaussian_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from prune import prune_list, calculate_v_imp_score +from torch.optim.lr_scheduler import ExponentialLR + +@register_method("lightgaussian") +class LightGaussianWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + + self.parser.add_argument("--prune_iterations", nargs="+", type=int, default=[16000, 24000]) + self.parser.add_argument("--prune_percent", type=float, default=0.5) + self.parser.add_argument("--v_pow", type=float, default=0.1) + self.parser.add_argument("--prune_decay", type=float, default=0.8) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + self.gaussians.scheduler = ExponentialLR(self.gaussians.optimizer, gamma=0.97) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.pruning_ratio_current = 0.0 + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + self.gaussians.scheduler.step() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + sti_value = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_t = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_p = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + if state is not None and "exp_avg_sq" in state: + vt = state["exp_avg_sq"] + eta_t = self.opt.position_lr_init + u_t = (eta_t / (torch.sqrt(vt) + 1e-8)) * grad_t + u_p = (eta_t / (torch.sqrt(vt) + 1e-8)) * grad_p + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + cos_sim = F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1) + grad_cos_sim = float(cos_sim.mean()) + norm_t = torch.norm(u_t[valid_mask], dim=1) + norm_p = torch.norm(u_p[valid_mask], dim=1) + ratio = norm_p / (norm_t + norm_p + 1e-7) + parasitic_ratio = float(ratio.mean()) + + conflict_mask = cos_sim < 0 + if conflict_mask.any(): + sti_value = float((ratio[conflict_mask] * torch.abs(cos_sim[conflict_mask])).mean()) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + prune_event = False + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step in self.args.prune_iterations: + prune_event = True + gaussian_list, imp_list = prune_list(self.gaussians, self.scene, self.pipe, self.background) + idx = self.args.prune_iterations.index(step) + v_list = calculate_v_imp_score(self.gaussians, imp_list, self.args.v_pow) + self.gaussians.prune_gaussians((self.args.prune_decay**idx) * self.args.prune_percent, v_list) + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + delta_n = int(num_gaussians - self.last_n_gaussians) + + if prune_event: + self.pruning_ratio_current = float(abs(delta_n) / (self.last_n_gaussians + 1e-7)) + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": delta_n, + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "prune_event": int(prune_event), "pruning_ratio_current": float(self.pruning_ratio_current) + } + + if self.track_decoupling: + metrics.update({ + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "sti": float(sti_value) + }) + + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + ply_path = os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply') + self.gaussians.load_ply(ply_path) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_lod_gs.py b/methods/wrapper_lod_gs.py new file mode 100644 index 0000000000000000000000000000000000000000..08b1e9757f7ad13ecefc7642e074305c6709230c --- /dev/null +++ b/methods/wrapper_lod_gs.py @@ -0,0 +1,266 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../LOD_GS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams, GMMParams + +@register_method("lod_gs") +class LODGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.gmmp = GMMParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.cap_gaussians = hyperparams.get("cap_gaussians", None) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gmm_args = self.gmmp.extract(self.args) + + self.pipe.antialiasing = False + + self.gaussians = GaussianModel(self.dataset.sh_degree, self.opt.optimizer_type, self.gmm_args) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + train_cams = self.scene.getTrainCameras() + self.viewpoint_stack = [c for item in train_cams for c in (item if isinstance(item, list) else [item])] + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + _cams = self.scene.getTrainCameras() + self.viewpoint_stack = [c for item in _cams for c in (item if isinstance(item, list) else [item])] + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + use_lod = step > 7000 + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, use_trained_exp=self.dataset.train_test_exp, LOD=use_lod) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + grad_t = {} + for name, param in [("xyz", self.gaussians._xyz), ("scale", self.gaussians._scaling), ("rot", self.gaussians._rotation), ("opac", self.gaussians._opacity), ("dc", self.gaussians._features_dc), ("rest", self.gaussians._features_rest), ("coefs", self.gaussians._coefs)]: + grad_t[name] = param.grad.clone() if param.grad is not None else torch.zeros_like(param) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + grad_p = {} + for name, param in [("xyz", self.gaussians._xyz), ("scale", self.gaussians._scaling), ("rot", self.gaussians._rotation), ("opac", self.gaussians._opacity), ("dc", self.gaussians._features_dc), ("rest", self.gaussians._features_rest), ("coefs", self.gaussians._coefs)]: + grad_p[name] = param.grad.clone() if param.grad is not None else torch.zeros_like(param) + + u_t_list = [] + u_p_list = [] + + name_mapping = {"xyz": "xyz", "scaling": "scale", "rotation": "rot", "opacity": "opac", "f_dc": "dc", "f_rest": "rest", "coefs": "coefs"} + + for group in self.gaussians.optimizer.param_groups: + g_name = group["name"] + if g_name in name_mapping: + p = group["params"][0] + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = group["lr"] + mapped_name = name_mapping[g_name] + + eff_t = (lr / (torch.sqrt(v_t) + 1e-7)) * grad_t[mapped_name] + eff_p = (lr / (torch.sqrt(v_t) + 1e-7)) * grad_p[mapped_name] + + u_t_list.append(eff_t.reshape(-1)) + u_p_list.append(eff_p.reshape(-1)) + + if len(u_t_list) > 0: + U_T = torch.cat(u_t_list) + U_P = torch.cat(u_p_list) + + mask = (torch.abs(U_P) > 1e-8) & (torch.abs(U_T) > 1e-8) + if mask.any(): + grad_cos_sim = float(F.cosine_similarity(U_T[mask].unsqueeze(0), U_P[mask].unsqueeze(0)).squeeze()) + parasitic_ratio = float(torch.norm(U_P) / (torch.norm(U_T) + torch.norm(U_P) + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + if self.cap_gaussians is None or self.gaussians.get_xyz.shape[0] < self.cap_gaussians: + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.015, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step < self.opt.iterations: + self.gaussians.exposure_optimizer.step() + self.gaussians.exposure_optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "lod_coefs_mag": float(torch.norm(self.gaussians._coefs, dim=-1).mean()) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, LOD=True) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + metrics["lod_coefs_mag"] = float(torch.norm(self.gaussians._coefs, dim=-1).mean()) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_lp3dgs.py b/methods/wrapper_lp3dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..1b4c47071d92a402529cc7ed1756818bc21f4d4e --- /dev/null +++ b/methods/wrapper_lp3dgs.py @@ -0,0 +1,270 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../LP-3DGS_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, ModelPool +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("lp3dgs") +class LP3DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.args.prune_method = "compact_3dgs" + self.args.use_importance_mask = True + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.prune_iterations = [19500] + self.opt.train_mask_iters = 500 + self.opt.gumbel_temp = 0.5 + self.opt.lambda_mask = 4e-3 + self.opt.mask_lr = 1e-3 + + if "360" in str(self.args.source_path): + self.opt.lambda_dist = 100.0 + else: + self.opt.lambda_dist = 1000.0 + + self.gaussians = ModelPool[self.dataset.prune_method](self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree(self.opt, step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + + mask_active = step > self.opt.prune_iterations[0] and step < self.opt.prune_iterations[0] + self.opt.train_mask_iters + if mask_active and ((step - self.opt.prune_iterations[0]) % 20 == 0): + self.gaussians.get_score_before_render(self.opt, scene=self.scene, pipe=self.pipe, background=bg, render=native_render) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + loss_parasitic = self.gaussians.addtional_loss(self.opt) if mask_active else torch.tensor(0.0, device="cuda") + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + sti_stats = {"sti_spatial": 0.0, "sti_geometry": 0.0, "sti_opacity": 0.0, "sti_appearance": 0.0} + + if self.track_decoupling and step % 100 == 0 and mask_active: + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + grad_t_flat = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in params if p.grad is not None]) + + grad_p_flat = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_t_flat, dim=1) > 0) & (torch.norm(grad_p_flat, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_t_flat[valid_mask], grad_p_flat[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_p_flat, dim=1).mean() / (torch.norm(grad_t_flat, dim=1).mean() + 1e-7)) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + sti_stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.gaussians.prune_before_render(self.opt, step) + + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, 1, scene=self.scene, pipe=self.pipe, background=bg, render=native_render, opt=self.opt) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.prune_after_render(self.opt, step, scene=self.scene, pipe=self.pipe, background=bg, render=native_render, use_importance_mask=self.args.use_importance_mask) + + if step < self.opt.iterations: + if mask_active: + self.gaussians.importance_mask.optimizer.step() + self.gaussians.importance_mask.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "loss_mask_penalty": float(loss_parasitic) + } + metrics.update(sti_stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_minisplatting.py b/methods/wrapper_minisplatting.py new file mode 100644 index 0000000000000000000000000000000000000000..4dd23ad6bc19836cc5603214657413f210ce1243 --- /dev/null +++ b/methods/wrapper_minisplatting.py @@ -0,0 +1,362 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../mini_splatting_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render_imp, render_depth +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from utils.sh_utils import SH2RGB + +@register_method("minisplatting") +class MiniSplattingWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.parser.add_argument("--simp_iteration1", type=int, default=15000) + self.parser.add_argument("--simp_iteration2", type=int, default=20000) + self.parser.add_argument("--num_depth", type=int, default=3500000) + self.parser.add_argument("--num_max", type=int, default=4500000) + self.parser.add_argument("--sampling_factor", type=float, default=0.5) + self.parser.add_argument("--imp_metric", type=str, default=None) + + self.args, _ = self.parser.parse_known_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + path_lower = str(self.args.source_path).lower() + if "360" in path_lower and not any(x in path_lower for x in ["room", "counter", "kitchen", "bonsai"]): + self.args.imp_metric = "outdoor" + else: + self.args.imp_metric = "indoor" + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.mask_blur = torch.zeros(self.last_n_gaussians, device="cuda", dtype=torch.bool) + + def train_iteration(self, step): + if step < self.args.simp_iteration1: + self.gaussians.update_learning_rate(step) + else: + self.gaussians.update_learning_rate(step - self.args.simp_iteration1 + 5000) + + if step % 1000 == 0 and step > self.args.simp_iteration1: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = render_imp(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + is_reinit_step = False + is_simplification_step = False + blur_mask_active_ratio = 0.0 + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + area_max = render_pkg.get("area_max", torch.zeros_like(visibility_filter, dtype=torch.float32)) + if area_max.shape[0] == self.gaussians._xyz.shape[0]: + condition = area_max > (image.shape[1] * image.shape[2] / 5000.0) + if self.mask_blur.shape[0] != self.gaussians._xyz.shape[0]: + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + self.mask_blur = torch.logical_or(self.mask_blur, condition) + + if self.mask_blur.shape[0] > 0: + blur_mask_active_ratio = float(self.mask_blur.float().mean()) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0 and step % 5000 != 0 and self.gaussians._xyz.shape[0] < self.args.num_max: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune_split(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, self.mask_blur) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 5000 == 0: + is_reinit_step = True + out_pts_list = [] + gt_list = [] + views = self.scene.getTrainCameras() + for view in views: + gt = view.original_image[0:3, :, :] + render_depth_pkg = render_depth(view, self.gaussians, self.pipe, self.background) + out_pts = render_depth_pkg["out_pts"] + accum_alpha = render_depth_pkg["accum_alpha"] + prob = 1.0 - accum_alpha + prob_sum = prob.sum() + if prob_sum > 0: + prob = prob / prob_sum + prob = prob.reshape(-1).cpu().numpy() + factor = 1.0 / (image.shape[1] * image.shape[2] * len(views) / self.args.num_depth) + N_xyz = prob.shape[0] + num_sampled = int(N_xyz * factor) + if num_sampled > 0 and num_sampled <= N_xyz: + indices = np.random.choice(N_xyz, size=num_sampled, p=prob, replace=False) + out_pts = out_pts.permute(1,2,0).reshape(-1,3) + gt = gt.permute(1,2,0).reshape(-1,3) + out_pts_list.append(out_pts[indices]) + gt_list.append(gt[indices]) + if out_pts_list: + out_pts_merged = torch.cat(out_pts_list) + gt_merged = torch.cat(gt_list) + self.gaussians.reinitial_pts(out_pts_merged, gt_merged) + self.gaussians.training_setup(self.opt) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + torch.cuda.empty_cache() + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + if step == self.args.simp_iteration1: + is_simplification_step = True + imp_score = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + accum_area_max = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + views = self.scene.getTrainCameras() + for view in views: + r_pkg = render_imp(view, self.gaussians, self.pipe, self.background) + accum_weights = r_pkg["accum_weights"] + area_proj = r_pkg["area_proj"] + area_max = r_pkg["area_max"] + accum_area_max = accum_area_max + area_max + if self.args.imp_metric == 'outdoor': + mask_t = area_max != 0 + temp = imp_score + accum_weights / (area_proj + 1e-7) + imp_score[mask_t] = temp[mask_t] + else: + imp_score = imp_score + accum_weights + imp_score[accum_area_max == 0] = 0 + imp_sum = imp_score.sum() + if imp_sum > 0: + prob = imp_score / imp_sum + else: + prob = imp_score + prob = prob.cpu().numpy() + factor = self.args.sampling_factor + N_xyz = self.gaussians._xyz.shape[0] + num_sampled = int(N_xyz * factor * ((prob != 0).sum() / prob.shape[0])) + if num_sampled > 0 and num_sampled <= N_xyz: + indices = np.random.choice(N_xyz, size=num_sampled, p=prob, replace=False) + mask = np.zeros(N_xyz, dtype=bool) + mask[indices] = True + self.gaussians.prune_points(~torch.tensor(mask, device="cuda")) + self.gaussians.max_sh_degree = self.dataset.sh_degree + self.gaussians.reinitial_pts(self.gaussians._xyz, SH2RGB(self.gaussians._features_dc + 0.0)[:,0]) + self.gaussians.training_setup(self.opt) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + torch.cuda.empty_cache() + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + if step == self.args.simp_iteration2: + is_simplification_step = True + imp_score = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + accum_area_max = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + views = self.scene.getTrainCameras() + for view in views: + r_pkg = render_imp(view, self.gaussians, self.pipe, self.background) + accum_weights = r_pkg["accum_weights"] + area_proj = r_pkg["area_proj"] + area_max = r_pkg["area_max"] + accum_area_max = accum_area_max + area_max + if self.args.imp_metric == 'outdoor': + mask_t = area_max != 0 + temp = imp_score + accum_weights / (area_proj + 1e-7) + imp_score[mask_t] = temp[mask_t] + else: + imp_score = imp_score + accum_weights + imp_score[accum_area_max == 0] = 0 + + vals, idx = torch.sort(imp_score + 1e-6) + cumsum_val = torch.cumsum(vals, dim=0) + split_indices = ((cumsum_val / vals.sum()) > (1.0 - 0.99)).nonzero() + if len(split_indices) > 0: + split_index = split_indices.min() + split_val_nonprune = vals[split_index] + non_prune_mask = imp_score > split_val_nonprune + else: + non_prune_mask = torch.ones_like(imp_score).bool() + + self.gaussians.prune_points(~non_prune_mask) + self.gaussians.training_setup(self.opt) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + torch.cuda.empty_cache() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "blur_mask_active_ratio": float(blur_mask_active_ratio), + "is_reinit_step": int(is_reinit_step), + "is_simplification_step": int(is_simplification_step) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = render_imp(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_minisplattingsgf.py b/methods/wrapper_minisplattingsgf.py new file mode 100644 index 0000000000000000000000000000000000000000..b880c6d6c10c6c6fdd92d5694f88fe9de44ba6fb --- /dev/null +++ b/methods/wrapper_minisplattingsgf.py @@ -0,0 +1,422 @@ +import os +import sys +import math +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../minisplatting_sgf_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render_imp, render_depth +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from utils.sh_utils import SH2RGB + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + +@register_method("minisplattingsgf") +class MiniSplattingSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.parser.add_argument("--simp_iteration1", type=int, default=15000) + self.parser.add_argument("--simp_iteration2", type=int, default=20000) + self.parser.add_argument("--num_depth", type=int, default=3500000) + self.parser.add_argument("--num_max", type=int, default=4500000) + self.parser.add_argument("--sampling_factor", type=float, default=0.5) + self.parser.add_argument("--imp_metric", type=str, default=None) + + self.args, _ = self.parser.parse_known_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + path_lower = str(self.args.source_path).lower() + if "360" in path_lower and not any(x in path_lower for x in ["room", "counter", "kitchen", "bonsai"]): + self.args.imp_metric = "outdoor" + else: + self.args.imp_metric = "indoor" + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.mask_blur = torch.zeros(self.last_n_gaussians, device="cuda", dtype=torch.bool) + + def train_iteration(self, step): + if step < self.args.simp_iteration1: + self.gaussians.update_learning_rate(step) + else: + self.gaussians.update_learning_rate(step - self.args.simp_iteration1 + 5000) + + if step % 1000 == 0 and step > self.args.simp_iteration1: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = render_imp(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_struct + loss_parasitic = loss_ssim_reweighted + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + is_reinit_step = False + is_simplification_step = False + blur_mask_active_ratio = 0.0 + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + area_max = render_pkg.get("area_max", torch.zeros_like(visibility_filter, dtype=torch.float32)) + if area_max.shape[0] == self.gaussians._xyz.shape[0]: + condition = area_max > (image.shape[1] * image.shape[2] / 5000.0) + if self.mask_blur.shape[0] != self.gaussians._xyz.shape[0]: + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + self.mask_blur = torch.logical_or(self.mask_blur, condition) + + if self.mask_blur.shape[0] > 0: + blur_mask_active_ratio = float(self.mask_blur.float().mean()) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0 and step % 5000 != 0 and self.gaussians._xyz.shape[0] < self.args.num_max: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune_split(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, self.mask_blur) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 5000 == 0: + is_reinit_step = True + out_pts_list = [] + gt_list = [] + views = self.scene.getTrainCameras() + for view in views: + gt = view.original_image[0:3, :, :] + render_depth_pkg = render_depth(view, self.gaussians, self.pipe, self.background) + out_pts = render_depth_pkg["out_pts"] + accum_alpha = render_depth_pkg["accum_alpha"] + prob = 1.0 - accum_alpha + prob_sum = prob.sum() + if prob_sum > 0: + prob = prob / prob_sum + prob = prob.reshape(-1).cpu().numpy() + factor = 1.0 / (image.shape[1] * image.shape[2] * len(views) / self.args.num_depth) + N_xyz = prob.shape[0] + num_sampled = int(N_xyz * factor) + if num_sampled > 0 and num_sampled <= N_xyz: + indices = np.random.choice(N_xyz, size=num_sampled, p=prob, replace=False) + out_pts = out_pts.permute(1,2,0).reshape(-1,3) + gt = gt.permute(1,2,0).reshape(-1,3) + out_pts_list.append(out_pts[indices]) + gt_list.append(gt[indices]) + if out_pts_list: + out_pts_merged = torch.cat(out_pts_list) + gt_merged = torch.cat(gt_list) + self.gaussians.reinitial_pts(out_pts_merged, gt_merged) + self.gaussians.training_setup(self.opt) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + torch.cuda.empty_cache() + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + if step == self.args.simp_iteration1: + is_simplification_step = True + imp_score = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + accum_area_max = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + views = self.scene.getTrainCameras() + for view in views: + r_pkg = render_imp(view, self.gaussians, self.pipe, self.background) + accum_weights = r_pkg["accum_weights"] + area_proj = r_pkg["area_proj"] + area_max = r_pkg["area_max"] + accum_area_max = accum_area_max + area_max + if self.args.imp_metric == 'outdoor': + mask_t = area_max != 0 + temp = imp_score + accum_weights / (area_proj + 1e-7) + imp_score[mask_t] = temp[mask_t] + else: + imp_score = imp_score + accum_weights + imp_score[accum_area_max == 0] = 0 + imp_sum = imp_score.sum() + if imp_sum > 0: + prob = imp_score / imp_sum + else: + prob = imp_score + prob = prob.cpu().numpy() + factor = self.args.sampling_factor + N_xyz = self.gaussians._xyz.shape[0] + num_sampled = int(N_xyz * factor * ((prob != 0).sum() / prob.shape[0])) + if num_sampled > 0 and num_sampled <= N_xyz: + indices = np.random.choice(N_xyz, size=num_sampled, p=prob, replace=False) + mask = np.zeros(N_xyz, dtype=bool) + mask[indices] = True + self.gaussians.prune_points(~torch.tensor(mask, device="cuda")) + self.gaussians.max_sh_degree = self.dataset.sh_degree + self.gaussians.reinitial_pts(self.gaussians._xyz, SH2RGB(self.gaussians._features_dc + 0.0)[:,0]) + self.gaussians.training_setup(self.opt) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + torch.cuda.empty_cache() + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + if step == self.args.simp_iteration2: + is_simplification_step = True + imp_score = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + accum_area_max = torch.zeros(self.gaussians._xyz.shape[0], device="cuda") + views = self.scene.getTrainCameras() + for view in views: + r_pkg = render_imp(view, self.gaussians, self.pipe, self.background) + accum_weights = r_pkg["accum_weights"] + area_proj = r_pkg["area_proj"] + area_max = r_pkg["area_max"] + accum_area_max = accum_area_max + area_max + if self.args.imp_metric == 'outdoor': + mask_t = area_max != 0 + temp = imp_score + accum_weights / (area_proj + 1e-7) + imp_score[mask_t] = temp[mask_t] + else: + imp_score = imp_score + accum_weights + imp_score[accum_area_max == 0] = 0 + + vals, idx = torch.sort(imp_score + 1e-6) + cumsum_val = torch.cumsum(vals, dim=0) + split_indices = ((cumsum_val / vals.sum()) > (1.0 - 0.99)).nonzero() + if len(split_indices) > 0: + split_index = split_indices.min() + split_val_nonprune = vals[split_index] + non_prune_mask = imp_score > split_val_nonprune + else: + non_prune_mask = torch.ones_like(imp_score).bool() + + self.gaussians.prune_points(~non_prune_mask) + self.gaussians.training_setup(self.opt) + self.mask_blur = torch.zeros(self.gaussians._xyz.shape[0], device="cuda", dtype=torch.bool) + torch.cuda.empty_cache() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "blur_mask_active_ratio": float(blur_mask_active_ratio), + "is_reinit_step": int(is_reinit_step), + "is_simplification_step": int(is_simplification_step), + "sgf_w_mean": float(W.mean()), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = render_imp(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_mipsplatting.py b/methods/wrapper_mipsplatting.py new file mode 100644 index 0000000000000000000000000000000000000000..746c26512eb3ff227b9fe6fc4ce3d91aa07ac390 --- /dev/null +++ b/methods/wrapper_mipsplatting.py @@ -0,0 +1,247 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../mip-splatting'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams, SGFParams +from utils.sgf_utils import SGFController + +@register_method("mipsplatting") +class MipSplattingWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.sp = SGFParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.sgf_args = self.sp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.train_cameras = self.scene.getTrainCameras().copy() + self.viewpoint_stack = self.train_cameras.copy() + + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + self.sgf_controller = SGFController(self.sgf_args) + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.train_cameras.copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + subpixel_offset = None + if self.dataset.ray_jitter: + subpixel_offset = torch.rand((int(viewpoint_cam.image_height), int(viewpoint_cam.image_width), 2), dtype=torch.float32, device="cuda") - 0.5 + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size, subpixel_offset=subpixel_offset) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + loss, struct_map, l1_map = self.sgf_controller.compute_loss(image, gt_image, lambda_dssim=self.opt.lambda_dssim) + + loss_target = (1.0 - self.opt.lambda_dssim) * l1_loss(image, gt_image) + loss_parasitic = loss - loss_target + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_t = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_p = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, {}) + exp_avg_sq = state.get('exp_avg_sq', torch.ones_like(self.gaussians._xyz)) + lr = self.opt.position_lr_init + for param_group in self.gaussians.optimizer.param_groups: + if param_group["name"] == "xyz": + lr = param_group['lr'] + + u_t = (lr / (torch.sqrt(exp_avg_sq) + 1e-8)) * grad_t + u_p = (lr / (torch.sqrt(exp_avg_sq) + 1e-8)) * grad_p + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0 and self.gaussians.get_xyz.shape[0] < 6_000_000: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter and step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(ssim(image, gt_image)), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "filter_3d_mean": float(self.gaussians.filter_3D.mean()) if hasattr(self.gaussians, "filter_3D") else 0.0, + "sgf_struct_map_mean": float(struct_map.mean()) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = self.gaussians.get_opacity_with_3D_filter.clone().detach() + scales = self.gaussians.get_scaling_with_3D_filter.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling_with_3D_filter + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_s, _ = torch.max(scales_2d, dim=1) + min_s, _ = torch.min(scales_2d, dim=1) + gamma = max_s / (min_s + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.gaussians.get_opacity_with_3D_filter)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + v = query_points.shape[0] + densities = torch.zeros(v, device="cuda") + xyz = self.gaussians._xyz + opacities = self.gaussians.get_opacity_with_3D_filter.squeeze() + scales = self.gaussians.get_scaling_with_3D_filter + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + n_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (n_gaussians + 1)) + for i in range(0, v, chunk_size): + end = min(i + chunk_size, v) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_mipsplattingsgf.py b/methods/wrapper_mipsplattingsgf.py new file mode 100644 index 0000000000000000000000000000000000000000..e61f0342c332ec814c8cdae0a03985b644a042a0 --- /dev/null +++ b/methods/wrapper_mipsplattingsgf.py @@ -0,0 +1,302 @@ +import os +import sys +import math +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../mip-splatting'))) +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + +@register_method("mipsplattingsgf") +class MipSplattingSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args, _ = self.parser.parse_known_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.train_cameras = self.scene.getTrainCameras().copy() + self.viewpoint_stack = self.train_cameras.copy() + + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.train_cameras.copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + subpixel_offset = None + if self.dataset.ray_jitter: + subpixel_offset = torch.rand((int(viewpoint_cam.image_height), int(viewpoint_cam.image_width), 2), dtype=torch.float32, device="cuda") - 0.5 + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size, subpixel_offset=subpixel_offset) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_struct + loss_parasitic = loss_ssim_reweighted + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_t = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_p = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, {}) + exp_avg_sq = state.get('exp_avg_sq', torch.ones_like(self.gaussians._xyz)) + lr = self.opt.position_lr_init + for param_group in self.gaussians.optimizer.param_groups: + if param_group["name"] == "xyz": + lr = param_group['lr'] + + u_t = (lr / (torch.sqrt(exp_avg_sq) + 1e-8)) * grad_t + u_p = (lr / (torch.sqrt(exp_avg_sq) + 1e-8)) * grad_p + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter and step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.train_cameras) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "filter_3d_mean": float(self.gaussians.filter_3D.mean()) if hasattr(self.gaussians, "filter_3D") else 0.0, + "sgf_w_mean": float(W.mean()), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = self.gaussians.get_opacity_with_3D_filter.clone().detach() + scales = self.gaussians.get_scaling_with_3D_filter.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling_with_3D_filter + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_s, _ = torch.max(scales_2d, dim=1) + min_s, _ = torch.min(scales_2d, dim=1) + gamma = max_s / (min_s + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(self.gaussians.get_opacity_with_3D_filter)) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + v = query_points.shape[0] + densities = torch.zeros(v, device="cuda") + xyz = self.gaussians._xyz + opacities = self.gaussians.get_opacity_with_3D_filter.squeeze() + scales = self.gaussians.get_scaling_with_3D_filter + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + n_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (n_gaussians + 1)) + for i in range(0, v, chunk_size): + end = min(i + chunk_size, v) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_ngs.py b/methods/wrapper_ngs.py new file mode 100644 index 0000000000000000000000000000000000000000..ad9ca48aa9ab8c85e9a7f016f711e2b76f4fd85c --- /dev/null +++ b/methods/wrapper_ngs.py @@ -0,0 +1,551 @@ +import os +import sys +import time +import math +import torch +import numpy as np +import torch.nn.functional as F + +from core.registry import register_method +from core.base_method import BaseMethod + +from ngs.config import TrainerConfig +from ngs.trainer import Runner + +@register_method("ngs") +class NGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.cfg = TrainerConfig() + self.cfg.data_dir = dataset_config["source_path"] + self.cfg.result_dir = dataset_config["model_path"] + self.cfg.data_factor = dataset_config.get("resolution", 1) + self.cfg.disable_viewer = True + + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.runner = Runner(0, 0, 1, self.cfg) + self.trainloader = torch.utils.data.DataLoader( + self.runner.trainset, + batch_size=self.cfg.batch_size, + shuffle=True, + num_workers=0, + persistent_workers=False, + pin_memory=True, + ) + self.trainloader_iter = iter(self.trainloader) + + self.schedulers = [ + torch.optim.lr_scheduler.ExponentialLR(self.runner.optimizers["means"], gamma=0.01 ** (1.0 / self.cfg.max_steps)), + ] + if self.cfg.pose_opt: + self.schedulers.append(torch.optim.lr_scheduler.ExponentialLR(self.runner.pose_optimizers[0], gamma=0.01 ** (1.0 / self.cfg.max_steps))) + if self.cfg.use_bilateral_grid: + self.schedulers.append( + torch.optim.lr_scheduler.ChainedScheduler([ + torch.optim.lr_scheduler.LinearLR(self.runner.bil_grid_optimizers[0], start_factor=0.01, total_iters=1000), + torch.optim.lr_scheduler.ExponentialLR(self.runner.bil_grid_optimizers[0], gamma=0.01 ** (1.0 / self.cfg.max_steps)), + ]) + ) + + self.last_n_gaussians = len(self.runner.splats["means"]) + self.freeze_surface = False + self.freeze_noise = False + self.global_tic = time.time() + + class _MockCamera: + def __init__(self, c2w, K, W, H): + if c2w.shape == (3, 4): + c2w = np.vstack([c2w, [[0, 0, 0, 1]]]) + w2c = np.linalg.inv(c2w) + self.world_view_transform = torch.tensor(w2c.T, dtype=torch.float32, device="cuda") + self.camera_center = torch.tensor(c2w[:3, 3], dtype=torch.float32, device="cuda") + fx, fy = float(K[0, 0]), float(K[1, 1]) + self.FoVx = 2.0 * math.atan(W / (2.0 * fx)) + self.FoVy = 2.0 * math.atan(H / (2.0 * fy)) + self.image_width = int(W) + self.image_height = int(H) + self.original_image = torch.zeros(3, int(H), int(W), device="cuda") + self.full_proj_transform = self.world_view_transform.clone() + + class _MockScene: + def __init__(self, runner): + self.runner = runner + + def _make_cams(self, indices): + parser = self.runner.parser + factor = int(getattr(parser, "factor", 1)) + cams = [] + for idx in indices: + i = int(idx) + c2w = np.asarray(parser.camtoworlds[i]) + cam_id = int(parser.camera_ids[i]) + K_orig = np.asarray(parser.Ks_dict[cam_id], dtype=np.float64).copy() + W_orig, H_orig = parser.imsize_dict[cam_id] + if factor > 1: + K = K_orig.copy() + K[0, 0] /= factor; K[1, 1] /= factor + K[0, 2] /= factor; K[1, 2] /= factor + W, H = W_orig // factor, H_orig // factor + else: + K, W, H = K_orig, W_orig, H_orig + cams.append(_MockCamera(c2w, K, W, H)) + return cams + + def getTrainCameras(self): + return self._make_cams(list(self.runner.trainset.indices)) + + def getTestCameras(self): + return self._make_cams(list(self.runner.valset.indices)) + + self.scene = _MockScene(self.runner) + + def train_iteration(self, step): + device = self.runner.device + cfg = self.cfg + + if step == cfg.noise_steps[0]: + self.freeze_surface = True + self.runner.initialize_and_prune_noise_gaussians( + erosion_steps=cfg.erosion_steps, increase_prune_steps=cfg.increase_erosion_steps + ) + from ngs.utils.noise import reset_opacity_noise_splats, reset_color_noise_splats + reset_opacity_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, opacity=cfg.init_noise_opa, device=device) + reset_color_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, noise_color=cfg.noise_color, device=device) + + if cfg.noise_steps[0] > 0 and step == cfg.noise_steps[0] + 1: + if cfg.reset_lr: + self.runner.optimizers["means"].param_groups[0]["lr"] = 1.6e-4 + self.schedulers[0] = torch.optim.lr_scheduler.ExponentialLR(self.runner.optimizers["means"], gamma=0.01 ** (1.0 / (cfg.max_steps - step))) + self.freeze_surface = False + self.freeze_noise = True + from ngs.utils.noise import reset_opacity_noise_splats, reset_color_noise_splats + reset_opacity_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, opacity=cfg.freeze_noise_opa, device=device) + reset_color_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, noise_color=cfg.noise_color, device=device) + + if len(cfg.noise_steps) == 2 and cfg.noise_steps[1] > 0 and step == cfg.noise_steps[1]: + from ngs.utils.noise import remove_noise_splats, reset_opacity_noise_splats, reset_color_noise_splats + remove_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, device=device) + self.runner.num_noise_gaussians = 0 + self.runner.noise_gaussian_data = None + self.runner.voxel_resolution = cfg.voxel_resolution + self.freeze_noise = False + self.freeze_surface = True + self.runner.initialize_and_prune_noise_gaussians( + erosion_steps=cfg.erosion_steps, increase_prune_steps=cfg.increase_erosion_steps + ) + reset_opacity_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, opacity=cfg.init_noise_opa, device=device) + reset_color_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, noise_color=cfg.noise_color, device=device) + + end_step = cfg.noise_steps[1] + cfg.unfreeze_steps[0] if len(cfg.noise_steps) == 2 and cfg.noise_steps[1] > 0 else cfg.max_steps + if step == end_step: + self.freeze_noise = True + self.freeze_surface = False + from ngs.utils.noise import reset_opacity_noise_splats, reset_color_noise_splats + reset_opacity_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, opacity=cfg.freeze_noise_opa, device=device) + reset_color_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, noise_color=cfg.noise_color, device=device) + + try: + data = next(self.trainloader_iter) + except StopIteration: + self.trainloader_iter = iter(self.trainloader) + data = next(self.trainloader_iter) + + camtoworlds = data["camtoworld"].to(device) + Ks = data["K"].to(device) + pixels = data["image"].to(device) + image_ids = data["image_id"].to(device) + masks = data["mask"].to(device) if "mask" in data else None + undist_masks = data["undist_mask"].to(device) if "undist_mask" in data else None + + if cfg.depth_loss: + points = data["points"].to(device) + depths_gt = data["depths"].to(device) + + height, width = pixels.shape[1:3] + + if cfg.pose_noise: + camtoworlds = self.runner.pose_perturb(camtoworlds, image_ids) + if cfg.pose_opt: + camtoworlds = self.runner.pose_adjust(camtoworlds, image_ids) + + sh_degree_to_use = min(step // cfg.sh_degree_interval, cfg.sh_degree) + + with torch.no_grad(): + if self.runner.filter_hull is None and step == cfg.hull_filter_step: + from ngs.utils.noise import compute_convex_hull_mesh + self.runner.filter_hull, _, _ = compute_convex_hull_mesh(self.runner.splats["means"], scale_factor=1.1) + if self.runner.num_noise_gaussians > 0 and cfg.noise_color is None: + from ngs.utils.noise import reset_color_noise_splats + reset_color_noise_splats(self.runner.splats, self.runner.optimizers, self.runner.num_noise_gaussians, device=device) + + renders, alphas, info = self.runner.rasterize_splats( + camtoworlds=camtoworlds, + Ks=Ks, + width=width, + height=height, + sh_degree=sh_degree_to_use, + near_plane=cfg.near_plane, + far_plane=cfg.far_plane, + image_ids=image_ids, + render_mode="RGB+ED" if cfg.depth_loss else "RGB", + undist_masks=undist_masks, + distributed=self.runner.world_size > 1, + ) + + if renders.shape[-1] == 4: + colors, depths = renders[..., 0:3], renders[..., 3:4] + else: + colors, depths = renders, None + + if cfg.use_bilateral_grid: + grid_y, grid_x = torch.meshgrid( + (torch.arange(height, device=device) + 0.5) / height, + (torch.arange(width, device=device) + 0.5) / width, + indexing="ij", + ) + grid_xy = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0) + from ngs.utils.lib_bilagrid import slice + colors = slice(self.runner.bil_grids, grid_xy, colors, image_ids)["rgb"] + + if cfg.random_bkgd: + color_choices = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 1.0, 0.0], [1.0, 0.0, 1.0], [0.0, 1.0, 1.0]], + device=device, + ) + random_indices = torch.randint(0, 6, (1, height, width), device=device) + bkgd = color_choices[random_indices] + if masks is not None: + bkgd[masks < 0.5] = 0.0 + colors = alphas * colors + bkgd * (1.0 - alphas) + + self.cfg.strategy.step_pre_backward( + params=self.runner.splats, + optimizers=self.runner.optimizers, + state=self.runner.strategy_state, + step=step, + info=info, + ) + + if masks is not None: + batch_indices, y_indices, x_indices = torch.where(masks > 0.1) + x_min, x_max = x_indices.min().item(), x_indices.max().item() + y_min, y_max = y_indices.min().item(), y_indices.max().item() + colors = colors[:, y_min:y_max, x_min:x_max] + pixels = pixels[:, y_min:y_max, x_min:x_max] + + l1loss = F.l1_loss(colors, pixels) + ssimloss = 1.0 - self.runner.ssim(colors.permute(0, 3, 1, 2), pixels.permute(0, 3, 1, 2)) + loss_target = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda + + loss_parasitic = torch.tensor(0.0, device=device) + + if cfg.depth_loss: + points_norm = torch.stack([points[:, :, 0] / (width - 1) * 2 - 1, points[:, :, 1] / (height - 1) * 2 - 1], dim=-1) + grid = points_norm.unsqueeze(2) + sampled_depths = F.grid_sample(depths.permute(0, 3, 1, 2), grid, align_corners=True).squeeze(3).squeeze(1) + disp = torch.where(sampled_depths > 0.0, 1.0 / sampled_depths, torch.zeros_like(sampled_depths)) + disp_gt = 1.0 / depths_gt + loss_parasitic += F.l1_loss(disp, disp_gt) * self.runner.scene_scale * cfg.depth_lambda + + if cfg.use_bilateral_grid: + from ngs.utils.lib_bilagrid import total_variation_loss + loss_parasitic += 10 * total_variation_loss(self.runner.bil_grids.grids) + + if masks is not None: + loss_parasitic += torch.sum(alphas * (1.0 - masks.unsqueeze(-1))) / ((1.0 - masks).sum()) + if cfg.foreground_loss: + from ngs.trainer import erode_masks + masks_eroded = erode_masks(masks, kernel_size=3, iterations=1) + loss_parasitic += torch.sum((1.0 - alphas) * masks_eroded.unsqueeze(-1)) / masks_eroded.sum() + + if cfg.opacity_reg > 0.0: + loss_parasitic += cfg.opacity_reg * torch.abs(torch.sigmoid(self.runner.splats["opacities"])).mean() + if cfg.scale_reg > 0.0: + loss_parasitic += cfg.scale_reg * torch.abs(torch.exp(self.runner.splats["scales"])).mean() + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats_sti = {} + + if self.track_decoupling and step % 100 == 0: + def get_effective_step(optimizer, param, grad): + state = optimizer.state.get(param, None) + lr = optimizer.param_groups[0]["lr"] + if state is not None and "exp_avg_sq" in state: + return (lr / (torch.sqrt(state["exp_avg_sq"]) + 1e-15)) * grad + return lr * grad + + params_appearance = [] + for k in ["sh0", "shN", "features", "colors"]: + if k in self.runner.splats: + params_appearance.append(self.runner.splats[k]) + + param_groups_map = { + "spatial": [self.runner.splats["means"]], + "geometry": [self.runner.splats["scales"], self.runner.splats["quats"]], + "opacity": [self.runner.splats["opacities"]], + "appearance": params_appearance, + } + + for opt in self.runner.optimizers.values(): + opt.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + grads_target = {} + for group_name, params_list in param_groups_map.items(): + vec = [] + for p in params_list: + if p.grad is not None: + opt = self.runner.optimizers.get([k for k, v in self.runner.splats.items() if v is p][0]) + step_vec = get_effective_step(opt, p, p.grad.clone()) + vec.append(step_vec.reshape(-1)) + grads_target[group_name] = torch.cat(vec) if vec else torch.tensor([], device=device) + + for opt in self.runner.optimizers.values(): + opt.zero_grad(set_to_none=True) + + if loss_parasitic.requires_grad: + loss_parasitic.backward(retain_graph=True) + + grads_parasitic = {} + for group_name, params_list in param_groups_map.items(): + vec = [] + for p in params_list: + if p.grad is not None: + opt = self.runner.optimizers.get([k for k, v in self.runner.splats.items() if v is p][0]) + step_vec = get_effective_step(opt, p, p.grad.clone()) + vec.append(step_vec.reshape(-1)) + grads_parasitic[group_name] = torch.cat(vec) if vec else torch.tensor([], device=device) + + valid_gt = [v for v in grads_target.values() if v.numel() > 0] + global_gt = torch.cat(valid_gt) if valid_gt else torch.tensor([0.0], device=device) + + valid_gp = [v for v in grads_parasitic.values() if v.numel() > 0] + global_gp = torch.cat(valid_gp) if valid_gp else torch.tensor([0.0], device=device) + + if global_gt.norm() > 0 and global_gp.norm() > 0: + grad_cos_sim = float(F.cosine_similarity(global_gt.unsqueeze(0), global_gp.unsqueeze(0))) + parasitic_ratio = float(global_gp.norm() / (global_gt.norm() + global_gp.norm() + 1e-7)) + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + ti = 0.0 + if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0: + if gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + stats_sti[f"sti_{group_name}"] = ti + + for opt in self.runner.optimizers.values(): + opt.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if "sh0" in self.runner.splats and self.runner.splats["sh0"].grad is not None: + self.runner.splats["sh0"].grad[: self.runner.num_noise_gaussians] = 0 + if "shN" in self.runner.splats and self.runner.splats["shN"].grad is not None: + self.runner.splats["shN"].grad[: self.runner.num_noise_gaussians] = 0 + if "colors" in self.runner.splats and self.runner.splats["colors"].grad is not None: + self.runner.splats["colors"].grad[: self.runner.num_noise_gaussians] = 0 + if "features" in self.runner.splats and self.runner.splats["features"].grad is not None: + self.runner.splats["features"].grad[: self.runner.num_noise_gaussians] = 0 + if "means" in self.runner.splats and self.runner.splats["means"].grad is not None: + self.runner.splats["means"].grad[: self.runner.num_noise_gaussians] = 0 + if "quats" in self.runner.splats and self.runner.splats["quats"].grad is not None: + self.runner.splats["quats"].grad[: self.runner.num_noise_gaussians] = 0 + if self.runner.splats["scales"].grad is not None: + self.runner.splats["scales"].grad[: self.runner.num_noise_gaussians] = 0 + if self.freeze_noise and "opacities" in self.runner.splats and self.runner.splats["opacities"].grad is not None: + self.runner.splats["opacities"].grad[: self.runner.num_noise_gaussians] = 0 + if self.freeze_surface: + for k in self.runner.splats.keys(): + if self.runner.splats[k].grad is not None: + self.runner.splats[k].grad[self.runner.num_noise_gaussians :] = 0 + + if cfg.sparse_grad and cfg.packed: + gaussian_ids = info["gaussian_ids"] + for k in self.runner.splats.keys(): + grad = self.runner.splats[k].grad + if grad is None or grad.is_sparse: + continue + self.runner.splats[k].grad = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=grad[gaussian_ids], + size=self.runner.splats[k].size(), + is_coalesced=len(Ks) == 1, + ) + + visibility_mask = None + if cfg.visible_adam: + if cfg.packed: + visibility_mask = torch.zeros_like(self.runner.splats["opacities"], dtype=torch.bool) + visibility_mask.scatter_(0, info["gaussian_ids"], 1) + else: + visibility_mask = (info["radii"] > 0).any(0) + + for optimizer in self.runner.optimizers.values(): + if cfg.visible_adam: + optimizer.step(visibility_mask) + else: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + for optimizer in self.runner.pose_optimizers: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + if self.freeze_surface: + for optimizer in self.runner.app_optimizers: + optimizer.zero_grad(set_to_none=True) + else: + for optimizer in self.runner.app_optimizers: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + for optimizer in self.runner.bil_grid_optimizers: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + + for scheduler in self.schedulers: + scheduler.step() + + self.runner.num_noise_gaussians, self.runner.noise_gaussian_data = self.cfg.strategy.step_post_backward( + params=self.runner.splats, + optimizers=self.runner.optimizers, + state=self.runner.strategy_state, + step=step, + info=info, + packed=cfg.packed, + num_noise=self.runner.num_noise_gaussians, + freeze_surface=self.freeze_surface, + noise_gaussian_data=self.runner.noise_gaussian_data, + filter_hull=self.runner.filter_hull, + ) + + num_gaussians = len(self.runner.splats["means"]) + metrics = { + "loss": float(loss), + "loss_l1": float(l1loss), + "loss_ssim": float(ssimloss), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "num_noise_gaussians": int(self.runner.num_noise_gaussians) + } + metrics.update(stats_sti) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.runner.splats["opacities"][self.runner.num_noise_gaussians:]).clone().detach() + scales = torch.exp(self.runner.splats["scales"][self.runner.num_noise_gaussians:]).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + histograms["anisotropy"] = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + if "sh0" in self.runner.splats: + histograms["sh_dc_mag"] = self.runner.splats["sh0"][self.runner.num_noise_gaussians:].detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + W = camera.image_width + H = camera.image_height + w2c = camera.world_view_transform.transpose(0, 1) + c2w = torch.linalg.inv(w2c) + + fovx = camera.FoVx + fovy = camera.FoVy + fx = W / (2.0 * math.tan(fovx / 2.0)) + fy = H / (2.0 * math.tan(fovy / 2.0)) + K = torch.tensor([[fx, 0.0, W / 2.0], [0.0, fy, H / 2.0], [0.0, 0.0, 1.0]], device="cuda") + + renders, _, _ = self.runner.rasterize_splats( + camtoworlds=c2w.unsqueeze(0), + Ks=K.unsqueeze(0), + width=W, + height=H, + sh_degree=self.cfg.sh_degree, + near_plane=self.cfg.near_plane, + far_plane=self.cfg.far_plane, + ) + + image = torch.clamp(renders[..., 0:3], 0.0, 1.0).squeeze(0).permute(2, 0, 1) + return {"image": image} + + def save(self, save_dir, step): + os.makedirs(os.path.join(save_dir, "ckpts"), exist_ok=True) + data = { + "step": step, + "splats": self.runner.splats.state_dict(), + "num_noise_gaussians": self.runner.num_noise_gaussians, + "noise_gaussian_data": self.runner.noise_gaussian_data + } + torch.save(data, os.path.join(save_dir, "ckpts", f"ckpt_{step}_rank0.pt")) + + def load(self, model_path, iteration): + ckpt_path = os.path.join(model_path, "ckpts", f"ckpt_{iteration}_rank0.pt") + if os.path.exists(ckpt_path): + ckpt = torch.load(ckpt_path, map_location=self.runner.device, weights_only=True) + for k in self.runner.splats.keys(): + self.runner.splats[k].data = ckpt["splats"][k] + + def get_spatial_centers(self): + return self.runner.splats["means"] + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.runner.splats["scales"][self.runner.num_noise_gaussians:]) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.runner.splats["opacities"][self.runner.num_noise_gaussians:]))) + + if "sh0" in self.runner.splats and "shN" in self.runner.splats: + dc = self.runner.splats["sh0"][self.runner.num_noise_gaussians:] + rest = self.runner.splats["shN"][self.runner.num_noise_gaussians:] + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = self.runner.splats["quats"][self.runner.num_noise_gaussians:].clone() + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.runner.splats["means"] + opacities = torch.sigmoid(self.runner.splats["opacities"]).squeeze() + scales = torch.exp(self.runner.splats["scales"]) + sigma_sq = scales[:, :2].max(dim=1)[0].pow(2) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + chunk_size = max(1, 30_000_000 // (xyz.shape[0] + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_octree_gs.py b/methods/wrapper_octree_gs.py new file mode 100644 index 0000000000000000000000000000000000000000..9420f8cf2c375572a013a5791e6aa8a9d408b737 --- /dev/null +++ b/methods/wrapper_octree_gs.py @@ -0,0 +1,250 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, "/root/autodl-tmp/Octree-GS_official") +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import prefilter_voxel, render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("octree_gs") +class OctreeGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel( + self.dataset.feat_dim, self.dataset.n_offsets, self.dataset.fork, self.dataset.use_feat_bank, self.dataset.appearance_dim, + self.dataset.add_opacity_dist, self.dataset.add_cov_dist, self.dataset.add_color_dist, self.dataset.add_level, + self.dataset.visible_threshold, self.dataset.dist2level, self.dataset.base_layer, self.dataset.progressive, self.dataset.extend + ) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + import logging; self.scene = Scene(self.dataset, self.gaussians, shuffle=False, resolution_scales=self.dataset.resolution_scales, logger=logging.getLogger("dummy")) + self.gaussians.training_setup(self.opt) + self.gaussians.set_coarse_interval(self.opt.coarse_iter, self.opt.coarse_factor) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = 0 + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + self.gaussians.set_anchor_mask(viewpoint_cam.camera_center, step, viewpoint_cam.resolution_scale) + voxel_visible_mask = prefilter_voxel(viewpoint_cam, self.gaussians, self.pipe, self.background) + retain_grad = (step < self.opt.update_until and step >= 0) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask, retain_grad=retain_grad) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + offset_selection_mask = render_pkg.get("selection_mask", torch.zeros_like(visibility_filter)) + scaling = render_pkg.get("scaling", torch.empty(0)) + opacity = render_pkg.get("neural_opacity", torch.empty(0)) + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + scaling_reg = scaling.prod(dim=1).mean() if scaling.shape[0] > 0 else torch.tensor(0.0, device="cuda") + loss_parasitic = 0.01 * scaling_reg + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and self.gaussians._anchor.grad is not None: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._anchor.grad.clone() if self.gaussians._anchor.grad is not None else torch.zeros_like(self.gaussians._anchor) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._anchor.grad.clone() if self.gaussians._anchor.grad is not None else torch.zeros_like(self.gaussians._anchor) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.update_until and step > self.opt.start_stat: + self.gaussians.training_statis(viewspace_point_tensor, opacity, visibility_filter, offset_selection_mask, voxel_visible_mask) + if self.opt.update_anchor and step > self.opt.update_from and step % self.opt.update_interval == 0: + self.gaussians.adjust_anchor( + iteration=step, + check_interval=self.opt.update_interval, + success_threshold=self.opt.success_threshold, + grad_threshold=self.opt.densify_grad_threshold, + update_ratio=self.dataset.update_ratio, + extra_ratio=self.dataset.extra_ratio, + extra_up=self.dataset.extra_up, + min_opacity=self.opt.min_opacity + ) + elif step == self.opt.update_until: + if hasattr(self.gaussians, 'opacity_accum'): + del self.gaussians.opacity_accum + del self.gaussians.offset_gradient_accum + del self.gaussians.offset_denom + torch.cuda.empty_cache() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = int(visibility_filter.sum().item()) if visibility_filter.shape[0] > 0 else 0 + num_anchors = self.gaussians.get_anchor.shape[0] + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(1.0 - ssim_value), + "num_gaussians": num_gaussians, "delta_N": num_gaussians - self.last_n_gaussians, + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "anchor_count": num_anchors, + "tree_level_mean": float(self.gaussians.get_level.float().mean()), + "spawn_ratio": float(num_gaussians / (num_anchors + 1e-7)) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = opacity.clone().detach() if opacity.shape[0] > 0 else torch.empty(0) + if scaling.shape[0] > 0: + histograms["scaling"] = scaling.clone().detach() + scales_2d = scaling[:, :2] if scaling.shape[1] >= 2 else scaling + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma.clone().detach() + histograms["sh_dc_mag"] = self.gaussians.get_anchor_feat.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + self.gaussians.set_anchor_mask(camera.camera_center, self.opt.iterations, camera.resolution_scale) + voxel_visible_mask = prefilter_voxel(camera, self.gaussians, self.pipe, self.background) + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply_sparse_gaussian(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians.get_anchor + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + if cameras is not None and len(cameras) > 0: + cam = cameras[0] + self.gaussians.set_anchor_mask(cam.camera_center, self.opt.iterations, cam.resolution_scale) + voxel_visible_mask = prefilter_voxel(cam, self.gaussians, self.pipe, self.background) + render_pkg = native_render(cam, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask) + + scales = render_pkg.get("scaling", torch.empty(0)) + opacities = render_pkg.get("neural_opacity", torch.empty(0)) + + if scales.shape[0] > 0: + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + + if opacities.shape[0] > 0: + metrics["alpha_mean"] = float(torch.mean(opacities)) + + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots_tensor = render_pkg.get("rotations", torch.empty(0)) + if rots_tensor.shape[0] > 0: + rots = F.normalize(rots_tensor.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians.get_anchor + + if xyz.shape[0] == 0: + return densities + + chunk_size = 500_000 + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + min_dist_sq, min_indices = torch.min(dist_sq, dim=1) + densities[i:end] = torch.exp(-0.5 * min_dist_sq) + return densities diff --git a/methods/wrapper_op43dgs.py b/methods/wrapper_op43dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..81464f4e8aa4c78c014a6b25c4ed8f5a55dc369f --- /dev/null +++ b/methods/wrapper_op43dgs.py @@ -0,0 +1,267 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/op43dgs_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("op43dgs") +class OP43DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + def get_eff_step(params): + eff_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p) + lr = 0.001 + for pg in self.gaussians.optimizer.param_groups: + if pg["params"][0] is p: + lr = pg["lr"] + break + if state is not None and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + step_val = (lr / (torch.sqrt(v) + 1e-15)) * p.grad.clone() + eff_grads.append(step_val.reshape(-1)) + else: + eff_grads.append((lr * p.grad.clone()).reshape(-1)) + return torch.cat(eff_grads) if eff_grads else None + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {k: get_eff_step(v) for k, v in param_groups_map.items()} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {k: get_eff_step(v) for k, v in param_groups_map.items()} + + for group_name in param_groups_map: + gt = grads_target.get(group_name) + gp = grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) + } + if self.track_decoupling and step % 100 == 0: + metrics["grad_cos_sim"] = float(grad_cos_sim) + metrics["parasitic_ratio"] = float(parasitic_ratio) + metrics.update(stats) + + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, "point_cloud", f"iteration_{iteration}", "point_cloud.ply")) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_opti3dgs.py b/methods/wrapper_opti3dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..a67443f89667deb8515fe6486dacc135d9e1830a --- /dev/null +++ b/methods/wrapper_opti3dgs.py @@ -0,0 +1,263 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Opti3DGS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("opti3dgs") +class Opti3DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.scaling_lr = 0.01 + self.opt.densification_interval = 500 + self.opt.densify_grad_threshold = 0.0001 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + k = 0 + p = 0 + if step < 12000: + if step < 1000: + k, p = 15, 7 + elif step < 4000: + k, p = 9, 4 + elif step < 8000: + k, p = 5, 2 + else: + k, p = 3, 1 + gt_image = F.avg_pool2d(gt_image, kernel_size=k, stride=1, padding=p) + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and step > 1: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_t = {} + for group in self.gaussians.optimizer.param_groups: + param = group['params'][0] + grads_t[group['name']] = param.grad.clone() if param.grad is not None else torch.zeros_like(param) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_p = {} + for group in self.gaussians.optimizer.param_groups: + param = group['params'][0] + grads_p[group['name']] = param.grad.clone() if param.grad is not None else torch.zeros_like(param) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + + cos_sims = [] + p_ratios = [] + for group in self.gaussians.optimizer.param_groups: + name = group['name'] + param = group['params'][0] + state = self.gaussians.optimizer.state.get(param, None) + if state is not None and 'exp_avg_sq' in state: + v = state['exp_avg_sq'] + lr = group['lr'] + u_t = (lr / (torch.sqrt(v) + 1e-8)) * grads_t[name] + u_p = (lr / (torch.sqrt(v) + 1e-8)) * grads_p[name] + + u_t_norm = torch.norm(u_t, dim=-1) + u_p_norm = torch.norm(u_p, dim=-1) + mask = (u_t_norm > 0) & (u_p_norm > 0) + + if mask.any(): + sim = F.cosine_similarity(u_t[mask], u_p[mask], dim=-1).mean() + cos_sims.append(sim) + r = u_p_norm[mask].mean() / (u_t_norm[mask].mean() + u_p_norm[mask].mean() + 1e-8) + p_ratios.append(r) + + if len(cos_sims) > 0: + grad_cos_sim = float(sum(cos_sims) / len(cos_sims)) + parasitic_ratio = float(sum(p_ratios) / len(p_ratios)) + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "blur_kernel_size": int(k), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_pgsr.py b/methods/wrapper_pgsr.py new file mode 100644 index 0000000000000000000000000000000000000000..13584acca07909df2aa322facc49e999dd485985 --- /dev/null +++ b/methods/wrapper_pgsr.py @@ -0,0 +1,395 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../pgsr_official'))) +from utils.loss_utils import l1_loss, ssim, lncc, get_img_grad_weight +from utils.graphics_utils import patch_offsets, patch_warp +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.app_model import AppModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from train import gen_virtul_cam + +@register_method("pgsr") +class PGSRWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + self.app_model = AppModel() + self.app_model.train() + self.app_model.cuda() + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + gt_image, gt_image_gray = viewpoint_cam.get_image() + + if step > 1000 and self.opt.exposure_compensation: + self.gaussians.use_app = True + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, app_model=self.app_model, + return_plane=step>self.opt.single_view_weight_from_iter, + return_depth_normal=step>self.opt.single_view_weight_from_iter) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + viewspace_point_tensor_abs = render_pkg.get("viewspace_points_abs", viewspace_point_tensor) + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + ssim_loss_val = (1.0 - ssim(image, gt_image)) + if 'app_image' in render_pkg and ssim_loss_val < 0.5: + app_image = render_pkg['app_image'] + Ll1 = l1_loss(app_image, gt_image) + else: + Ll1 = l1_loss(image, gt_image) + + image_loss = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * ssim_loss_val + loss = image_loss.clone() + + loss_target = image_loss.clone() + loss_parasitic = torch.tensor(0.0, device="cuda") + + if visibility_filter.sum() > 0: + scale = self.gaussians.get_scaling[visibility_filter] + sorted_scale, _ = torch.sort(scale, dim=-1) + min_scale_loss = sorted_scale[...,0] + scale_l = self.opt.scale_loss_weight * min_scale_loss.mean() + loss += scale_l + loss_parasitic += scale_l + + normal_loss_val = 0.0 + if step > self.opt.single_view_weight_from_iter: + weight = self.opt.single_view_weight + normal = render_pkg["rendered_normal"] + depth_normal = render_pkg["depth_normal"] + image_weight = (1.0 - get_img_grad_weight(gt_image)) + image_weight = (image_weight).clamp(0,1).detach() ** 2 + if not self.opt.wo_image_weight: + normal_loss = weight * (image_weight * (((depth_normal - normal)).abs().sum(0))).mean() + else: + normal_loss = weight * (((depth_normal - normal)).abs().sum(0)).mean() + loss += normal_loss + loss_parasitic += normal_loss + normal_loss_val = float(normal_loss) + + geo_loss_val = 0.0 + ncc_loss_val = 0.0 + if step > self.opt.multi_view_weight_from_iter: + nearest_id = getattr(viewpoint_cam, 'nearest_id', []) + nearest_cam = None if len(nearest_id) == 0 else self.scene.getTrainCameras()[random.sample(nearest_id, 1)[0]] + use_virtul_cam = False + if self.opt.use_virtul_cam and (np.random.random() < self.opt.virtul_cam_prob or nearest_cam is None): + nearest_cam = gen_virtul_cam(viewpoint_cam, trans_noise=self.dataset.multi_view_max_dis, deg_noise=self.dataset.multi_view_max_angle) + use_virtul_cam = True + + if nearest_cam is not None: + try: + patch_size = self.opt.multi_view_patch_size + sample_num = self.opt.multi_view_sample_num + pixel_noise_th = self.opt.multi_view_pixel_noise_th + total_patch_size = (patch_size * 2 + 1) ** 2 + ncc_weight = self.opt.multi_view_ncc_weight + geo_weight = self.opt.multi_view_geo_weight + + H, W = render_pkg['plane_depth'].squeeze().shape + ix, iy = torch.meshgrid(torch.arange(W), torch.arange(H), indexing='xy') + pixels = torch.stack([ix, iy], dim=-1).float().to(render_pkg['plane_depth'].device) + + nearest_render_pkg = native_render(nearest_cam, self.gaussians, self.pipe, bg, app_model=self.app_model, return_plane=True, return_depth_normal=False) + + pts = self.gaussians.get_points_from_depth(viewpoint_cam, render_pkg['plane_depth']) + pts_in_nearest_cam = pts @ nearest_cam.world_view_transform[:3,:3] + nearest_cam.world_view_transform[3,:3] + map_z, d_mask = self.gaussians.get_points_depth_in_depth_map(nearest_cam, nearest_render_pkg['plane_depth'], pts_in_nearest_cam) + + pts_in_nearest_cam = pts_in_nearest_cam / (pts_in_nearest_cam[:,2:3]) + pts_in_nearest_cam = pts_in_nearest_cam * map_z.squeeze()[...,None] + R = torch.tensor(nearest_cam.R).float().cuda() + T = torch.tensor(nearest_cam.T).float().cuda() + pts_ = (pts_in_nearest_cam-T)@R.transpose(-1,-2) + pts_in_view_cam = pts_ @ viewpoint_cam.world_view_transform[:3,:3] + viewpoint_cam.world_view_transform[3,:3] + pts_projections = torch.stack([ + pts_in_view_cam[:,0] * viewpoint_cam.Fx / pts_in_view_cam[:,2] + viewpoint_cam.Cx, + pts_in_view_cam[:,1] * viewpoint_cam.Fy / pts_in_view_cam[:,2] + viewpoint_cam.Cy + ], -1).float() + + pixel_noise = torch.norm(pts_projections - pixels.reshape(*pts_projections.shape), dim=-1) + if not self.opt.wo_use_geo_occ_aware: + d_mask = d_mask & (pixel_noise < pixel_noise_th) + weights = (1.0 / torch.exp(pixel_noise)).detach() + weights[~d_mask] = 0 + else: + weights = torch.ones_like(pixel_noise) + weights[~d_mask] = 0 + + if d_mask.sum() > 0: + geo_loss = geo_weight * ((weights * pixel_noise)[d_mask]).mean() + loss += geo_loss + loss_parasitic += geo_loss + geo_loss_val = float(geo_loss) + + if not use_virtul_cam: + with torch.no_grad(): + d_mask = d_mask.reshape(-1) + valid_indices = torch.arange(d_mask.shape[0], device=d_mask.device)[d_mask] + if d_mask.sum() > sample_num: + index = np.random.choice(d_mask.sum().cpu().numpy(), sample_num, replace=False) + valid_indices = valid_indices[index] + weights = weights.reshape(-1)[valid_indices] + pixels = pixels.reshape(-1,2)[valid_indices] + offsets = patch_offsets(patch_size, pixels.device) + ori_pixels_patch = pixels.reshape(-1, 1, 2) / viewpoint_cam.ncc_scale + offsets.float() + H, W = gt_image_gray.squeeze().shape + pixels_patch = ori_pixels_patch.clone() + pixels_patch[:, :, 0] = 2 * pixels_patch[:, :, 0] / (W - 1) - 1.0 + pixels_patch[:, :, 1] = 2 * pixels_patch[:, :, 1] / (H - 1) - 1.0 + ref_gray_val = F.grid_sample(gt_image_gray.unsqueeze(1), pixels_patch.view(1, -1, 1, 2), align_corners=True) + ref_gray_val = ref_gray_val.reshape(-1, total_patch_size) + + ref_to_neareast_r = nearest_cam.world_view_transform[:3,:3].transpose(-1,-2) @ viewpoint_cam.world_view_transform[:3,:3] + ref_to_neareast_t = -ref_to_neareast_r @ viewpoint_cam.world_view_transform[3,:3] + nearest_cam.world_view_transform[3,:3] + + ref_local_n = render_pkg["rendered_normal"].permute(1,2,0).reshape(-1,3)[valid_indices] + ref_local_d = render_pkg['rendered_distance'].squeeze().reshape(-1)[valid_indices] + + H_ref_to_neareast = ref_to_neareast_r[None] - torch.matmul(ref_to_neareast_t[None,:,None].expand(ref_local_d.shape[0],3,1), ref_local_n[:,:,None].expand(ref_local_d.shape[0],3,1).permute(0, 2, 1))/ref_local_d[...,None,None] + + K_nearest = nearest_cam.get_k(nearest_cam.ncc_scale) + if isinstance(K_nearest, np.ndarray): K_nearest = torch.from_numpy(K_nearest) + K_nearest = K_nearest.to(torch.float32).cuda() + + K_viewpoint_inv = viewpoint_cam.get_inv_k(viewpoint_cam.ncc_scale) + if isinstance(K_viewpoint_inv, np.ndarray): K_viewpoint_inv = torch.from_numpy(K_viewpoint_inv) + K_viewpoint_inv = K_viewpoint_inv.to(torch.float32).cuda() + + H_ref_to_neareast = torch.matmul(K_nearest[None].expand(ref_local_d.shape[0], 3, 3), H_ref_to_neareast) + H_ref_to_neareast = H_ref_to_neareast @ K_viewpoint_inv + + grid = patch_warp(H_ref_to_neareast.reshape(-1,3,3), ori_pixels_patch) + grid[:, :, 0] = 2 * grid[:, :, 0] / (W - 1) - 1.0 + grid[:, :, 1] = 2 * grid[:, :, 1] / (H - 1) - 1.0 + _, nearest_image_gray = nearest_cam.get_image() + sampled_gray_val = F.grid_sample(nearest_image_gray[None], grid.reshape(1, -1, 1, 2), align_corners=True) + sampled_gray_val = sampled_gray_val.reshape(-1, total_patch_size) + + ncc, ncc_mask = lncc(ref_gray_val, sampled_gray_val) + mask = ncc_mask.reshape(-1) + ncc = ncc.reshape(-1) * weights + ncc = ncc[mask].squeeze() + + if mask.sum() > 0: + ncc_loss = ncc_weight * ncc.mean() + loss += ncc_loss + loss_parasitic += ncc_loss + ncc_loss_val = float(ncc_loss) + except Exception: + pass + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + obs_mask = (render_pkg.get("out_observe", torch.ones_like(visibility_filter)) > 0) & visibility_filter + self.gaussians.max_radii2D[obs_mask] = torch.max(self.gaussians.max_radii2D[obs_mask], radii[obs_mask]) + self.gaussians.add_densification_stats(viewspace_point_tensor, viewspace_point_tensor_abs, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.densify_abs_grad_threshold, self.opt.opacity_cull_threshold, self.scene.cameras_extent, size_threshold) + + if self.opt.use_multi_view_trim and step % 1000 == 0 and step < self.opt.densify_until_iter: + observe_the = 2 + observe_cnt = torch.zeros_like(self.gaussians.get_opacity) + for view in self.scene.getTrainCameras(): + render_pkg_tmp = native_render(view, self.gaussians, self.pipe, bg, app_model=self.app_model, return_plane=False, return_depth_normal=False) + out_observe = render_pkg_tmp["out_observe"] + observe_cnt[out_observe > 0] += 1 + prune_mask = (observe_cnt < observe_the).squeeze() + if prune_mask.sum() > 0: + self.gaussians.prune_points(prune_mask) + + if step < self.opt.densify_until_iter: + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.app_model.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(ssim_loss_val), + "loss_normal": float(normal_loss_val), "loss_geo": float(geo_loss_val), "loss_ncc": float(ncc_loss_val), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, app_model=self.app_model, return_plane=False, return_depth_normal=False) + return {"image": render_pkg.get("app_image", render_pkg["render"]), "depth": render_pkg.get("plane_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + self.app_model.save_weights(save_dir, step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + self.app_model.load_weights(model_path, iteration) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_pgsrsgf.py b/methods/wrapper_pgsrsgf.py new file mode 100644 index 0000000000000000000000000000000000000000..2ee45d11274d7cdb30bd31d4faa2901a6c62a144 --- /dev/null +++ b/methods/wrapper_pgsrsgf.py @@ -0,0 +1,473 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../pgsr_official'))) +from utils.loss_utils import l1_loss, ssim, lncc, get_img_grad_weight +from utils.graphics_utils import patch_offsets, patch_warp +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from scene.app_model import AppModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from train import gen_virtul_cam + + +# === SGF helper functions === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("pgsrsgf") +class PGSRSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + # PGSR-specific: appearance model + self.app_model = AppModel() + self.app_model.train() + self.app_model.cuda() + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + gt_image, gt_image_gray = viewpoint_cam.get_image() + + # PGSR-specific: exposure compensation + if step > 1000 and self.opt.exposure_compensation: + self.gaussians.use_app = True + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, app_model=self.app_model, + return_plane=step > self.opt.single_view_weight_from_iter, + return_depth_normal=step > self.opt.single_view_weight_from_iter) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + viewspace_point_tensor_abs = render_pkg.get("viewspace_points_abs", viewspace_point_tensor) + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + # === SGF: structure-aware reweighted photometric loss === + # Use app_image for L1 when available (consistent with PGSR design), + # but always use raw `image` for Sobel structure detection. + sgf_l1_image = image + if 'app_image' in render_pkg: + ssim_check = (1.0 - ssim(image, gt_image)) + if ssim_check < 0.5: + sgf_l1_image = render_pkg['app_image'] + + E_struct = torch.zeros(image.shape[1], image.shape[2], device=image.device) + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W_sgf = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(sgf_l1_image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W_sgf)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W_sgf)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + # === PGSR-specific: scale loss === + if visibility_filter.sum() > 0: + scale = self.gaussians.get_scaling[visibility_filter] + sorted_scale, _ = torch.sort(scale, dim=-1) + min_scale_loss = sorted_scale[..., 0] + scale_l = self.opt.scale_loss_weight * min_scale_loss.mean() + loss += scale_l + loss_parasitic += scale_l + + # === PGSR-specific: single-view normal consistency === + normal_loss_val = 0.0 + if step > self.opt.single_view_weight_from_iter: + weight = self.opt.single_view_weight + normal = render_pkg["rendered_normal"] + depth_normal = render_pkg["depth_normal"] + image_weight = (1.0 - get_img_grad_weight(gt_image)) + image_weight = (image_weight).clamp(0, 1).detach() ** 2 + if not self.opt.wo_image_weight: + normal_loss = weight * (image_weight * (((depth_normal - normal)).abs().sum(0))).mean() + else: + normal_loss = weight * (((depth_normal - normal)).abs().sum(0)).mean() + loss += normal_loss + loss_parasitic += normal_loss + normal_loss_val = float(normal_loss) + + # === PGSR-specific: multi-view geometric + NCC losses === + geo_loss_val = 0.0 + ncc_loss_val = 0.0 + if step > self.opt.multi_view_weight_from_iter: + nearest_id = getattr(viewpoint_cam, 'nearest_id', []) + nearest_cam = None if len(nearest_id) == 0 else self.scene.getTrainCameras()[random.sample(nearest_id, 1)[0]] + use_virtul_cam = False + if self.opt.use_virtul_cam and (np.random.random() < self.opt.virtul_cam_prob or nearest_cam is None): + nearest_cam = gen_virtul_cam(viewpoint_cam, trans_noise=self.dataset.multi_view_max_dis, deg_noise=self.dataset.multi_view_max_angle) + use_virtul_cam = True + + if nearest_cam is not None: + try: + patch_size = self.opt.multi_view_patch_size + sample_num = self.opt.multi_view_sample_num + pixel_noise_th = self.opt.multi_view_pixel_noise_th + total_patch_size = (patch_size * 2 + 1) ** 2 + ncc_weight = self.opt.multi_view_ncc_weight + geo_weight = self.opt.multi_view_geo_weight + + H, W = render_pkg['plane_depth'].squeeze().shape + ix, iy = torch.meshgrid(torch.arange(W), torch.arange(H), indexing='xy') + pixels = torch.stack([ix, iy], dim=-1).float().to(render_pkg['plane_depth'].device) + + nearest_render_pkg = native_render(nearest_cam, self.gaussians, self.pipe, bg, app_model=self.app_model, return_plane=True, return_depth_normal=False) + + pts = self.gaussians.get_points_from_depth(viewpoint_cam, render_pkg['plane_depth']) + pts_in_nearest_cam = pts @ nearest_cam.world_view_transform[:3, :3] + nearest_cam.world_view_transform[3, :3] + map_z, d_mask = self.gaussians.get_points_depth_in_depth_map(nearest_cam, nearest_render_pkg['plane_depth'], pts_in_nearest_cam) + + pts_in_nearest_cam = pts_in_nearest_cam / (pts_in_nearest_cam[:, 2:3]) + pts_in_nearest_cam = pts_in_nearest_cam * map_z.squeeze()[..., None] + R = torch.tensor(nearest_cam.R).float().cuda() + T = torch.tensor(nearest_cam.T).float().cuda() + pts_ = (pts_in_nearest_cam - T) @ R.transpose(-1, -2) + pts_in_view_cam = pts_ @ viewpoint_cam.world_view_transform[:3, :3] + viewpoint_cam.world_view_transform[3, :3] + pts_projections = torch.stack([ + pts_in_view_cam[:, 0] * viewpoint_cam.Fx / pts_in_view_cam[:, 2] + viewpoint_cam.Cx, + pts_in_view_cam[:, 1] * viewpoint_cam.Fy / pts_in_view_cam[:, 2] + viewpoint_cam.Cy + ], -1).float() + + pixel_noise = torch.norm(pts_projections - pixels.reshape(*pts_projections.shape), dim=-1) + if not self.opt.wo_use_geo_occ_aware: + d_mask = d_mask & (pixel_noise < pixel_noise_th) + weights = (1.0 / torch.exp(pixel_noise)).detach() + weights[~d_mask] = 0 + else: + weights = torch.ones_like(pixel_noise) + weights[~d_mask] = 0 + + if d_mask.sum() > 0: + geo_loss = geo_weight * ((weights * pixel_noise)[d_mask]).mean() + loss += geo_loss + loss_parasitic += geo_loss + geo_loss_val = float(geo_loss) + + if not use_virtul_cam: + with torch.no_grad(): + d_mask = d_mask.reshape(-1) + valid_indices = torch.arange(d_mask.shape[0], device=d_mask.device)[d_mask] + if d_mask.sum() > sample_num: + index = np.random.choice(d_mask.sum().cpu().numpy(), sample_num, replace=False) + valid_indices = valid_indices[index] + weights = weights.reshape(-1)[valid_indices] + pixels = pixels.reshape(-1, 2)[valid_indices] + offsets = patch_offsets(patch_size, pixels.device) + ori_pixels_patch = pixels.reshape(-1, 1, 2) / viewpoint_cam.ncc_scale + offsets.float() + H, W = gt_image_gray.squeeze().shape + pixels_patch = ori_pixels_patch.clone() + pixels_patch[:, :, 0] = 2 * pixels_patch[:, :, 0] / (W - 1) - 1.0 + pixels_patch[:, :, 1] = 2 * pixels_patch[:, :, 1] / (H - 1) - 1.0 + ref_gray_val = F.grid_sample(gt_image_gray.unsqueeze(1), pixels_patch.view(1, -1, 1, 2), align_corners=True) + ref_gray_val = ref_gray_val.reshape(-1, total_patch_size) + + ref_to_neareast_r = nearest_cam.world_view_transform[:3, :3].transpose(-1, -2) @ viewpoint_cam.world_view_transform[:3, :3] + ref_to_neareast_t = -ref_to_neareast_r @ viewpoint_cam.world_view_transform[3, :3] + nearest_cam.world_view_transform[3, :3] + + ref_local_n = render_pkg["rendered_normal"].permute(1, 2, 0).reshape(-1, 3)[valid_indices] + ref_local_d = render_pkg['rendered_distance'].squeeze().reshape(-1)[valid_indices] + + H_ref_to_neareast = ref_to_neareast_r[None] - torch.matmul( + ref_to_neareast_t[None, :, None].expand(ref_local_d.shape[0], 3, 1), + ref_local_n[:, :, None].expand(ref_local_d.shape[0], 3, 1).permute(0, 2, 1) + ) / ref_local_d[..., None, None] + + K_nearest = nearest_cam.get_k(nearest_cam.ncc_scale) + if isinstance(K_nearest, np.ndarray): K_nearest = torch.from_numpy(K_nearest) + K_nearest = K_nearest.to(torch.float32).cuda() + + K_viewpoint_inv = viewpoint_cam.get_inv_k(viewpoint_cam.ncc_scale) + if isinstance(K_viewpoint_inv, np.ndarray): K_viewpoint_inv = torch.from_numpy(K_viewpoint_inv) + K_viewpoint_inv = K_viewpoint_inv.to(torch.float32).cuda() + + H_ref_to_neareast = torch.matmul(K_nearest[None].expand(ref_local_d.shape[0], 3, 3), H_ref_to_neareast) + H_ref_to_neareast = H_ref_to_neareast @ K_viewpoint_inv + + grid = patch_warp(H_ref_to_neareast.reshape(-1, 3, 3), ori_pixels_patch) + grid[:, :, 0] = 2 * grid[:, :, 0] / (W - 1) - 1.0 + grid[:, :, 1] = 2 * grid[:, :, 1] / (H - 1) - 1.0 + _, nearest_image_gray = nearest_cam.get_image() + sampled_gray_val = F.grid_sample(nearest_image_gray[None], grid.reshape(1, -1, 1, 2), align_corners=True) + sampled_gray_val = sampled_gray_val.reshape(-1, total_patch_size) + + ncc, ncc_mask = lncc(ref_gray_val, sampled_gray_val) + mask = ncc_mask.reshape(-1) + ncc = ncc.reshape(-1) * weights + ncc = ncc[mask].squeeze() + + if mask.sum() > 0: + ncc_loss = ncc_weight * ncc.mean() + loss += ncc_loss + loss_parasitic += ncc_loss + ncc_loss_val = float(ncc_loss) + except Exception: + pass + + # === Gradient decoupling tracking === + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + # === Densification (PGSR backbone: 3-arg stats, 5-arg prune) === + with torch.no_grad(): + if step < self.opt.densify_until_iter: + obs_mask = (render_pkg.get("out_observe", torch.ones_like(visibility_filter)) > 0) & visibility_filter + self.gaussians.max_radii2D[obs_mask] = torch.max(self.gaussians.max_radii2D[obs_mask], radii[obs_mask]) + self.gaussians.add_densification_stats(viewspace_point_tensor, viewspace_point_tensor_abs, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, self.opt.densify_abs_grad_threshold, self.opt.opacity_cull_threshold, self.scene.cameras_extent, size_threshold) + + # PGSR-specific: multi-view trim + if self.opt.use_multi_view_trim and step % 1000 == 0 and step < self.opt.densify_until_iter: + observe_the = 2 + observe_cnt = torch.zeros_like(self.gaussians.get_opacity) + for view in self.scene.getTrainCameras(): + render_pkg_tmp = native_render(view, self.gaussians, self.pipe, bg, app_model=self.app_model, return_plane=False, return_depth_normal=False) + out_observe = render_pkg_tmp["out_observe"] + observe_cnt[out_observe > 0] += 1 + prune_mask = (observe_cnt < observe_the).squeeze() + if prune_mask.sum() > 0: + self.gaussians.prune_points(prune_mask) + + if step < self.opt.densify_until_iter: + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.app_model.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.app_model.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "loss_normal": float(normal_loss_val), "loss_geo": float(geo_loss_val), "loss_ncc": float(ncc_loss_val), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W_sgf.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted), + "iter_time_ms": float((time.perf_counter() - _iter_start) * 1000), + "vram_allocated_GB": float(torch.cuda.memory_allocated() / 1024**3), + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, app_model=self.app_model, return_plane=False, return_depth_normal=False) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg.get("app_image", render_pkg["render"]), "depth": render_pkg.get("plane_depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + self.app_model.save_weights(save_dir, step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + self.app_model.load_weights(model_path, iteration) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_pixelgs.py b/methods/wrapper_pixelgs.py new file mode 100644 index 0000000000000000000000000000000000000000..217ddde704da136ec046b29773e355ed4ee05e72 --- /dev/null +++ b/methods/wrapper_pixelgs.py @@ -0,0 +1,196 @@ +import os +import sys +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Pixel-GS'))) +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from utils.loss_utils import l1_loss, ssim + +@register_method("pixelgs") +class PixelGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, depth_threshold=self.opt.depth_threshold * self.scene.cameras_extent) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + pixels = render_pkg["pixels"] + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter, pixels) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_pixelgssgf.py b/methods/wrapper_pixelgssgf.py new file mode 100644 index 0000000000000000000000000000000000000000..50b8c0f4a820d68608ee580eccd3cbad6736f0f0 --- /dev/null +++ b/methods/wrapper_pixelgssgf.py @@ -0,0 +1,262 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Pixel-GS'))) +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + +@register_method("pixelgssgf") +class PixelGSSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args, _ = self.parser.parse_known_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, depth_threshold=self.opt.depth_threshold * self.scene.cameras_extent) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + pixels = render_pkg["pixels"] + gt_image = viewpoint_cam.original_image.cuda() + + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter, pixels) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_pup3dgs.py b/methods/wrapper_pup3dgs.py new file mode 100644 index 0000000000000000000000000000000000000000..98780e023f59fc2e459ee9b9c321287660178247 --- /dev/null +++ b/methods/wrapper_pup3dgs.py @@ -0,0 +1,268 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/PUP-3DGS_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("pup3dgs") +class PUP3DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + loss_ssim = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - loss_ssim) + loss = loss_target + loss_parasitic + + stats = {} + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + stats["grad_cos_sim"] = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + else: + stats["grad_cos_sim"] = 0.0 + stats["parasitic_ratio"] = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest] + } + + def extract_effective_step(): + grads = {} + for g_name, params in param_groups_map.items(): + g_list = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = next((g["lr"] for g in self.gaussians.optimizer.param_groups if any(p is param for param in g["params"])), 1.0) + eff = (lr / (torch.sqrt(v_t) + 1e-15)) * p.grad.clone() + g_list.append(eff.reshape(-1)) + else: + g_list.append(p.grad.clone().reshape(-1)) + if g_list: + grads[g_name] = torch.cat(g_list) + else: + grads[g_name] = torch.tensor([], device="cuda") + return grads + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = extract_effective_step() + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = extract_effective_step() + + for g_name in param_groups_map: + gt = grads_target[g_name] + gp = grads_parasitic[g_name] + if gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{g_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_s, _ = torch.max(scales_2d, dim=1) + min_s, _ = torch.min(scales_2d, dim=1) + gamma = max_s / (min_s + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + v_size = query_points.shape[0] + densities = torch.zeros(v_size, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + n_pts = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (n_pts + 1)) + for i in range(0, v_size, chunk_size): + end = min(i + chunk_size, v_size) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_qgs.py b/methods/wrapper_qgs.py new file mode 100644 index 0000000000000000000000000000000000000000..493b3b3d8483384614f80910134a54f905251365 --- /dev/null +++ b/methods/wrapper_qgs.py @@ -0,0 +1,298 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../quadratic_gs_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("qgs") +class QGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args, _ = self.parser.parse_known_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.densify_grad_threshold = 0.3 + self.opt.percent_dense = 0.001 + + self.dataset.root_dir = self.dataset.source_path + self.dataset.dataset = self.dataset + self.dataset.gs_model = self.dataset + self.dataset.optimizer = self.opt + self.dataset.pipeline = self.pipe + self.dataset.downsample = self.args.resolution + self.dataset.undistortion = False + self.dataset.use_alpha = False + self.dataset.type = "Colmap" + self.dataset.hold = 8 + self.dataset.multi_view_num = 5; self.dataset.multi_view_max_angle = 15.0; self.dataset.multi_view_max_dis = 1.0; self.dataset.multi_view_min_dis = 0.0 + self.pipe.depth_ratio = 0.0; self.pipe.reciprocal_z = False + self.pipe.occlusion_awared_denom = False + self.opt.occupancy_rate = 0.0 + + self.gaussians = GaussianModel(self.dataset) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + return_depth = step > self.opt.distortion_from_iter + return_normal = step > self.opt.depth_normal_from_iter + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.opt, self.background, kernel_size=self.dataset.kernel_size, stop_z_gradient=False, return_depth=return_depth, return_normal=return_normal) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + loss_parasitic = torch.tensor(0.0, device="cuda") + + if step > self.opt.distortion_from_iter and "render_dist" in render_pkg: + loss_parasitic += self.opt.lambda_dist * render_pkg["render_dist"].mean() + + if step > self.opt.depth_normal_from_iter and "render_normal" in render_pkg and "surf_normal" in render_pkg: + normal = render_pkg["render_normal"] + depth_normal = render_pkg["surf_normal"] + if "render_curvature" in render_pkg: + rend_curvature_log = torch.clamp_max(torch.log(torch.clamp_min(torch.nan_to_num(render_pkg["render_curvature"].squeeze(), 0.0001), 0.0001)), 1.0) + image_weight = 1.0 - torch.sigmoid(rend_curvature_log) + normal_loss = self.opt.lambda_normal * (image_weight.detach() * (1 - (normal * depth_normal)).sum(dim=0)[None]).mean() + loss_parasitic += normal_loss + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and loss_parasitic > 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + target_grads = {} + for group in self.gaussians.optimizer.param_groups: + p = group['params'][0] + if p.grad is not None: + target_grads[group['name']] = p.grad.clone() + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + parasitic_grads = {} + for group in self.gaussians.optimizer.param_groups: + p = group['params'][0] + if p.grad is not None: + parasitic_grads[group['name']] = p.grad.clone() + + sims = [] + p_norms = [] + t_norms = [] + + for group in self.gaussians.optimizer.param_groups: + name = group['name'] + if name in target_grads and name in parasitic_grads: + state = self.gaussians.optimizer.state.get(group['params'][0], None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = group["lr"] + + gt = target_grads[name] + gp = parasitic_grads[name] + + ut = (lr / (torch.sqrt(v_t) + 1e-15)) * gt + up = (lr / (torch.sqrt(v_t) + 1e-15)) * gp + + ut_flat = ut.view(ut.shape[0], -1) + up_flat = up.view(up.shape[0], -1) + + mask = (torch.norm(ut_flat, dim=1) > 0) & (torch.norm(up_flat, dim=1) > 0) + if mask.any(): + sim = F.cosine_similarity(ut_flat[mask], up_flat[mask], dim=1).mean() + sims.append(sim) + + t_norms.append(torch.norm(ut_flat, dim=1).mean()) + p_norms.append(torch.norm(up_flat, dim=1).mean()) + + if sims: + grad_cos_sim = float(torch.stack(sims).mean()) + if t_norms and p_norms: + parasitic_ratio = float(torch.stack(p_norms).mean() / (torch.stack(t_norms).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, radii, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling[:, :3]).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.opt, self.background, kernel_size=self.dataset.kernel_size, stop_z_gradient=True, return_depth=True, return_normal=True) + return {"image": render_pkg["render"], "depth": render_pkg.get("surf_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales[:, :3]) + shape_params = torch.tanh(raw_scales[:, 3:6]) + + scales_2d = scales[:, :2] + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + s3_vals = torch.abs(shape_params[:, 2]) + metrics["curvature_s3_median"] = float(torch.median(s3_vals)) + metrics["quadric_deformation_ratio"] = float((s3_vals > 0.1).float().mean()) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling[:, :3]) + sigma_sq = scales[:, :2].max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_r3dg.py b/methods/wrapper_r3dg.py new file mode 100644 index 0000000000000000000000000000000000000000..10211f300f7c6dff753750518b81940bc0aa2769 --- /dev/null +++ b/methods/wrapper_r3dg.py @@ -0,0 +1,250 @@ +import os +import sys +import random +import torch +import math +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Relightable3DGaussian'))) +from utils.loss_utils import ssim, first_order_edge_aware_loss +from gaussian_renderer.render import render as native_render +from gaussian_renderer.render import render_view +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("r3dg") +class R3DGWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.type = "render" + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree, render_type=self.args.type) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = render_view(viewpoint_cam, self.gaussians, self.pipe, self.background, 1.0, None, computer_pseudo_normal=True) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + weights = render_pkg["weights"] + + gt_image = viewpoint_cam.original_image.cuda() + image_mask = viewpoint_cam.image_mask.cuda() if hasattr(viewpoint_cam, "image_mask") and viewpoint_cam.image_mask is not None else torch.ones_like(gt_image[:1]) + + Ll1 = F.l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + loss_parasitic = torch.tensor(0.0, device="cuda") + + if self.opt.lambda_mask_entropy > 0: + o = render_pkg["opacity"].clamp(1e-6, 1 - 1e-6) + loss_parasitic += self.opt.lambda_mask_entropy * -(image_mask * torch.log(o) + (1-image_mask) * torch.log(1 - o)).mean() + + if self.opt.lambda_normal_render_depth > 0: + loss_parasitic += self.opt.lambda_normal_render_depth * F.mse_loss(render_pkg["normal"] * image_mask, render_pkg["pseudo_normal"].detach() * image_mask) + + if self.opt.lambda_normal_smooth > 0: + loss_parasitic += self.opt.lambda_normal_smooth * first_order_edge_aware_loss(render_pkg["normal"], gt_image) + + if self.opt.lambda_depth_smooth > 0: + loss_parasitic += self.opt.lambda_depth_smooth * first_order_edge_aware_loss(render_pkg["depth"], gt_image) + + if self.opt.lambda_depth_var > 0: + lambda_depth_var = self.opt.lambda_depth_var * min(math.pow(10, step / 5000), 100) + loss_parasitic += lambda_depth_var * render_pkg["depth_var"].clamp_min(1e-6).sqrt().mean() + + if self.opt.lambda_scaling > 0: + scaling = self.gaussians.get_scaling + scaling_loss = (scaling - scaling.mean(dim=-1, keepdim=True)).abs().sum(-1).mean() + lambda_scaling = self.opt.lambda_scaling - 0.99 * self.opt.lambda_scaling * min(1, 4 * step / self.opt.iterations) + loss_parasitic += lambda_scaling * scaling_loss + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if loss_parasitic.requires_grad: + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter, weights) + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + densify_grad_normal_threshold = self.opt.densify_grad_normal_threshold if step > self.opt.normal_densify_from_iter else 99999 + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, densify_grad_normal_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._shs_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, opt=self.opt) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._shs_dc, self.gaussians._shs_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_radegs.py b/methods/wrapper_radegs.py new file mode 100644 index 0000000000000000000000000000000000000000..f7c1d3bed9414517abd357cf5d2e5eae38f2903a --- /dev/null +++ b/methods/wrapper_radegs.py @@ -0,0 +1,367 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +target_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../RaDe-GS_offy')) +if not os.path.exists(target_path): + raise FileNotFoundError(f"Target directory {target_path} does not exist.") + +if target_path not in sys.path: + sys.path.insert(0, target_path) + +for key in list(sys.modules.keys()): + if key.split('.')[0] in ['utils', 'scene', 'gaussian_renderer', 'arguments', 'warp_patch_ncc', 'fused_ssim', 'simple_knn']: + del sys.modules[key] + +from utils.loss_utils import ssim, L1_loss_appearance, PatchMatch +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from utils.graphics_utils import depth_double_to_normal, depth_to_normal + +@register_method("radegs") +class RaDeGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.kernel_size = self.dataset.kernel_size + + self.trainCameras = self.scene.getTrainCameras().copy() + if self.dataset.disable_filter3D: + self.gaussians.reset_3D_filter() + else: + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + self.patchmatch = None + if self.opt.lambda_multi_view_ncc > 0 or self.opt.lambda_multi_view_geo > 0: + self.patchmatch = PatchMatch( + self.opt.multi_view_patch_size, + self.opt.multi_view_pixel_noise_th, + kernel_size=self.kernel_size, + pipe=self.pipe, + debug=False, + model_path=self.dataset.model_path, + ) + + self.viewpoint_stack = None + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + reg_kick_on = step >= self.opt.regularization_from_iter + + render_pkg = native_render( + viewpoint_cam, + self.gaussians, + self.pipe, + self.background, + self.kernel_size, + require_depth=reg_kick_on + ) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = L1_loss_appearance(image, gt_image, self.gaussians, viewpoint_cam.uid) + ssim_value = ssim(image.unsqueeze(0), gt_image.unsqueeze(0)) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + depth_normal_loss = torch.tensor([0], dtype=torch.float32, device="cuda") + ncc_loss = torch.tensor([0], dtype=torch.float32, device="cuda") + geo_loss = torch.tensor([0], dtype=torch.float32, device="cuda") + depth_normal = None + + if reg_kick_on and self.opt.lambda_depth_normal > 0: + rendered_expected_depth = render_pkg["expected_depth"] + rendered_median_depth = render_pkg["median_depth"] + rendered_normal = render_pkg["normal"] + if 0.0 < self.dataset.depth_ratio < 1.0: + depth_normal = depth_double_to_normal(viewpoint_cam, rendered_expected_depth, rendered_median_depth) + normal_error_map = 1 - torch.linalg.vecdot(rendered_normal.unsqueeze(0), depth_normal, dim=1) + depth_normal_loss = (1 - self.dataset.depth_ratio) * normal_error_map[0].mean() + self.dataset.depth_ratio * normal_error_map[1].mean() + depth_normal = None + else: + depth_map = rendered_expected_depth if self.dataset.depth_ratio < 1.0 else rendered_median_depth + depth_normal = depth_to_normal(viewpoint_cam, depth_map) + normal_error_map = 1 - torch.linalg.vecdot(rendered_normal, depth_normal, dim=0) + depth_normal_loss = normal_error_map.mean() + + if reg_kick_on and self.opt.lambda_multi_view_ncc > 0 and self.patchmatch is not None: + nearest_cam = None if len(viewpoint_cam.nearest_id) == 0 else self.scene.getTrainCameras()[random.sample(viewpoint_cam.nearest_id, 1)[0]] + ncc_loss, geo_loss = self.patchmatch(self.gaussians, render_pkg, viewpoint_cam, nearest_cam, step, depth_normal) + + loss_parasitic = self.opt.lambda_depth_normal * depth_normal_loss + self.opt.lambda_multi_view_ncc * ncc_loss + self.opt.lambda_multi_view_geo * geo_loss + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + sti_stats = {"sti_spatial": 0.0, "sti_geometry": 0.0, "sti_opacity": 0.0, "sti_appearance": 0.0} + + if self.track_decoupling and step % 100 == 0: + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_t_global = [] + grads_target = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + group_grads.append(p.grad.clone().reshape(-1)) + grad_t_global.append(p.grad.clone().reshape(-1)) + if group_grads: + grads_target[group_name] = torch.cat(group_grads) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if hasattr(loss_parasitic, 'requires_grad') and loss_parasitic.requires_grad: + loss_parasitic.backward(retain_graph=True) + grad_p_global = [] + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + group_grads.append(p.grad.clone().reshape(-1)) + grad_p_global.append(p.grad.clone().reshape(-1)) + if group_grads: + grads_parasitic[group_name] = torch.cat(group_grads) + + if grad_t_global and grad_p_global: + vec_t = torch.cat(grad_t_global) + vec_p = torch.cat(grad_p_global) + norm_t = vec_t.norm() + norm_p = vec_p.norm() + if norm_t > 0 and norm_p > 0: + grad_cos_sim = float(F.cosine_similarity(vec_t.unsqueeze(0), vec_p.unsqueeze(0))) + parasitic_ratio = float(norm_p / (norm_t + norm_p + 1e-7)) + + for group_name in param_groups_map: + gt = grads_target.get(group_name) + gp = grads_parasitic.get(group_name) + if gt is not None and gp is not None: + lr_val = 0.001 + for pg in self.gaussians.optimizer.param_groups: + if pg["name"] == group_name or (group_name == "spatial" and pg["name"] == "xyz") or (group_name == "geometry" and pg["name"] in ["scaling", "rotation"]): + lr_val = pg["lr"] + break + + eps = 1e-15 + v_t_list = [] + v_p_list = [] + for p in param_groups_map[group_name]: + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + v_t_list.append(v.reshape(-1)) + v_p_list.append(v.reshape(-1)) + else: + v_t_list.append(torch.ones_like(p).reshape(-1)) + v_p_list.append(torch.ones_like(p).reshape(-1)) + + v_t_tensor = torch.cat(v_t_list) + v_p_tensor = torch.cat(v_p_list) + + u_t = (lr_val / (torch.sqrt(v_t_tensor) + eps)) * gt + u_p = (lr_val / (torch.sqrt(v_p_tensor) + eps)) * gp + + norm_u_t = u_t.norm() + norm_u_p = u_p.norm() + + if norm_u_t > 0 and norm_u_p > 0: + cos_val = float(F.cosine_similarity(u_t.unsqueeze(0), u_p.unsqueeze(0))) + r_val = float(norm_u_p / (norm_u_t + norm_u_p + 1e-7)) + ti = r_val * max(0.0, -cos_val) + else: + ti = 0.0 + sti_stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.05, self.scene.cameras_extent, size_threshold) + if self.dataset.disable_filter3D: + self.gaussians.reset_3D_filter() + else: + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % 100 == 0 and step > self.opt.densify_until_iter and not self.dataset.disable_filter3D: + if step < self.opt.iterations - 100: + self.gaussians.compute_3D_filter(cameras=self.trainCameras) + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "ncc_loss": float(ncc_loss), + "depth_normal_loss": float(depth_normal_loss) + } + metrics.update(sti_stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = self.gaussians.get_scaling.clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, self.kernel_size) + return {"image": render_pkg["render"], "depth": render_pkg.get("expected_depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, "point_cloud", f"iteration_{iteration}", "point_cloud.ply")) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = self.gaussians.get_opacity_with_3D_filter.squeeze() + scales = self.gaussians.get_scaling_with_3D_filter + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_rdogaussian.py b/methods/wrapper_rdogaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..77f15ff27c4e28c302a74e9f0b4c31de8eefd6d9 --- /dev/null +++ b/methods/wrapper_rdogaussian.py @@ -0,0 +1,357 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../RDO-Gaussian_official'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from vq.ecvq import get_vq_cfg + +@register_method("rdogaussian") +class RDOGaussianWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.vq_cfg = get_vq_cfg(self.opt) + self.gaussians = GaussianModel(self.dataset.sh_degree, vq_cfg=self.vq_cfg) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + + self.gaussians.initialize_quantizer() + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.activate_vq = False + self.activate_gsprune = False + self.activate_shprune = False + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + + if step == self.opt.vq_start_iter: + for key in self.vq_cfg["keys"]: + self.gaussians._quantizer[key].kmeans_initialize(self.gaussians.vector_dict[key].data) + self.activate_vq = True + + if self.activate_vq and step == self.opt.rate_constrain_iter: + for key in self.vq_cfg["keys"]: + self.gaussians._quantizer[key].rate_constrain = True + + if step == self.opt.gs_prune_start_iter: + self.activate_gsprune = True + + if step == self.opt.sh_prune_start_iter: + self.activate_shprune = True + + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render( + viewpoint_cam, self.gaussians, self.pipe, self.background, + vq_cfg=self.vq_cfg, + activate_vq=self.activate_vq, + activate_shprune=self.activate_shprune, + activate_gsprune=self.activate_gsprune, + update_index=(step - 1) % 10 == 0 + ) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + rate_loss = render_pkg["rate_loss"] + vq_loss = render_pkg["vq_loss"] + sh_mask_loss = render_pkg["sh_mask_loss"] + sh_mask_percent = render_pkg["sh_mask_percent"] + gs_mask_loss = render_pkg["gs_mask_loss"] + gs_mask_percent = render_pkg["gs_mask_percent"] + bits_dict = render_pkg["bits"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + loss_parasitic = vq_loss + rate_loss + self.opt.sh_mask_lambda * sh_mask_loss + self.opt.gs_mask_lambda * gs_mask_loss + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and self.activate_gsprune: + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_gs_mask.zero_grad(set_to_none=True) + self.gaussians.optimizer_sh_mask.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + grad_T_spatial = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_gs_mask.zero_grad(set_to_none=True) + self.gaussians.optimizer_sh_mask.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + grad_P_spatial = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_T_spatial, dim=1) > 0) & (torch.norm(grad_P_spatial, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_T_spatial[valid_mask], grad_P_spatial[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_P_spatial, dim=1).mean() / (torch.norm(grad_T_spatial, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.optimizer_gs_mask.zero_grad(set_to_none=True) + self.gaussians.optimizer_sh_mask.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if self.activate_vq and (step % self.opt.reactivate_codeword_period == 0 and step < self.opt.rate_constrain_iter): + for key in self.vq_cfg["keys"]: + self.gaussians._quantizer[key].reactivate_codeword(prob_threshold=1e-5) + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + if self.activate_vq: + self.gaussians.optimizer_vq.step() + self.gaussians.scheduler_vq.step() + self.gaussians.optimizer_vq.zero_grad(set_to_none=True) + if self.activate_shprune: + self.gaussians.optimizer_sh_mask.step() + self.gaussians.optimizer_sh_mask.zero_grad(set_to_none=True) + if self.activate_gsprune: + self.gaussians.optimizer_gs_mask.step() + self.gaussians.optimizer_gs_mask.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_target - Ll1), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "vq_loss": float(vq_loss), "rate_loss": float(rate_loss), + "gs_mask_percent": float(gs_mask_percent), "sh_mask_percent": float(sh_mask_percent), + "total_bits": float(sum(bits_dict.values())) if bits_dict else 0.0, + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + with torch.no_grad(): + if self.activate_gsprune: + gs_prune_dict = self.gaussians.apply_gs_mask() + scales = gs_prune_dict['scale'] + opacities = gs_prune_dict['opacity'] + valid_mask = gs_prune_dict['gs_bitmask'] + else: + scales = self.gaussians.get_scaling + opacities = self.gaussians.get_opacity + valid_mask = torch.ones(scales.shape[0], dtype=torch.bool, device="cuda") + + valid_opacities = opacities[valid_mask] + valid_scales = scales[valid_mask] + + histograms["opacity"] = valid_opacities.clone().detach() + histograms["scaling"] = valid_scales.clone().detach() + scales_2d = valid_scales[:, :2] if valid_scales.shape[1] >= 2 else valid_scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + + if self.activate_shprune: + sh_prune_dict = self.gaussians.apply_sh_mask() + f_dc = self.gaussians._features_dc[valid_mask] + else: + f_dc = self.gaussians._features_dc[valid_mask] + histograms["sh_dc_mag"] = f_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render( + camera, self.gaussians, self.pipe, self.background, + vq_cfg=self.vq_cfg, + activate_vq=self.activate_vq, + activate_shprune=self.activate_shprune, + activate_gsprune=self.activate_gsprune, + update_index=False + ) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + bitstream_path = os.path.join(model_path, "bitstream") + ply_path = os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply') + if os.path.exists(bitstream_path): + self.gaussians.decode(bitstream_path) + self.activate_vq = True + self.activate_gsprune = True + self.activate_shprune = True + else: + self.gaussians.load_ply(ply_path) + + def get_spatial_centers(self): + with torch.no_grad(): + if self.activate_gsprune: + gs_prune_dict = self.gaussians.apply_gs_mask() + valid_mask = gs_prune_dict['gs_bitmask'] + return self.gaussians._xyz[valid_mask] + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + if self.activate_gsprune: + gs_prune_dict = self.gaussians.apply_gs_mask() + scales = gs_prune_dict['scale'] + opacities = gs_prune_dict['opacity'] + valid_mask = gs_prune_dict['gs_bitmask'] + else: + scales = self.gaussians.get_scaling + opacities = self.gaussians.get_opacity + valid_mask = torch.ones(scales.shape[0], dtype=torch.bool, device="cuda") + + valid_scales = scales[valid_mask] + valid_opacities = opacities[valid_mask] + + scales_2d = valid_scales[:, :2] if valid_scales.dim() > 1 and valid_scales.shape[1] >= 2 else valid_scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(valid_opacities)) + + if self.activate_shprune: + sh_prune_dict = self.gaussians.apply_sh_mask() + rest = sh_prune_dict['features_rest'][valid_mask] + else: + rest = self.gaussians._features_rest[valid_mask] + + dc = self.gaussians._features_dc[valid_mask] + + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = self.gaussians.get_rotation[valid_mask] + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + + if self.activate_gsprune: + gs_prune_dict = self.gaussians.apply_gs_mask() + scales = gs_prune_dict['scale'] + opacities = gs_prune_dict['opacity'].squeeze() + valid_mask = gs_prune_dict['gs_bitmask'] + else: + scales = self.gaussians.get_scaling + opacities = self.gaussians.get_opacity.squeeze() + valid_mask = torch.ones(scales.shape[0], dtype=torch.bool, device="cuda") + + xyz = self.gaussians._xyz[valid_mask] + opacities = opacities[valid_mask] + scales = scales[valid_mask] + + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + if N_gaussians == 0: + return densities + + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + diff --git a/methods/wrapper_reactgs.py b/methods/wrapper_reactgs.py new file mode 100644 index 0000000000000000000000000000000000000000..a9bb20fca8878a4f9e1cac10ccf998b86367476e --- /dev/null +++ b/methods/wrapper_reactgs.py @@ -0,0 +1,253 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../ReAct-GS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("reactgs") +class ReActGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.densify_grad_threshold = 0.0003 + self.opt.nsp_interval = 3000 + self.opt.nsp_util = 20000 + self.opt.densify_from_iter = 500 + self.opt.densify_until_iter = 15000 + self.opt.densification_interval = 100 + self.opt.depth_threshold = 0.37 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg, depth_threshold=self.opt.depth_threshold * self.scene.cameras_extent) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + importance = render_pkg.get("weights", torch.zeros_like(radii)) / (render_pkg.get("pixels", torch.ones_like(radii)) + 1e-8) + importance[importance.isnan()] = 0 + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = 0.0 + for param_group in self.gaussians.optimizer.param_groups: + if param_group["name"] == "xyz": + lr = param_group['lr'] + break + u_target = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_parasitic = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_target, dim=1) > 0) & (torch.norm(u_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target[valid_mask], u_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_parasitic, dim=1).mean() / (torch.norm(u_target, dim=1).mean() + torch.norm(u_parasitic, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification(viewspace_point_tensor, visibility_filter, importance) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step % self.opt.nsp_interval == 0 and step < self.opt.nsp_util: + self.gaussians.needle_shape_perturbance() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + + with torch.no_grad(): + scales_eval = torch.exp(self.gaussians._scaling) + max_scales_eval, _ = torch.max(scales_eval, dim=1) + sum_scales_eval = torch.sum(scales_eval, dim=1) + needle_ratio_tensor = max_scales_eval / (sum_scales_eval + 1e-7) + needle_gaussian_ratio = float(((needle_ratio_tensor > 0.8) & (needle_ratio_tensor < 0.999)).float().mean()) + mean_importance = float(importance.mean()) + + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "mean_importance": mean_importance, "needle_gaussian_ratio": needle_gaussian_ratio + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, depth_threshold=self.opt.depth_threshold * self.scene.cameras_extent) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_scaffoldgs.py b/methods/wrapper_scaffoldgs.py new file mode 100644 index 0000000000000000000000000000000000000000..f7f138b14add7475ad3fc53494d2882feb0ce010 --- /dev/null +++ b/methods/wrapper_scaffoldgs.py @@ -0,0 +1,236 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +_UPSTREAM = '/root/autodl-tmp/scaffold_gs_official' +sys.path.insert(0, _UPSTREAM) +# monorepo-flat: insert each first-level subdir individually +import os as _os +for _sub in ('utils', 'scene', 'gaussian_renderer', 'arguments', 'lpipsPyTorch', 'experiments'): + _p = _os.path.join(_UPSTREAM, _sub) + if _os.path.isdir(_p): + sys.path.insert(0, _p) +del _os, _p, _sub, _UPSTREAM +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("scaffoldgs") +class ScaffoldGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_anchors = self.gaussians.get_anchor.shape[0] if hasattr(self.gaussians, 'get_anchor') else 0 + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + pass + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - getattr(self.opt, 'lambda_dssim', 0.2)) * Ll1 + loss_parasitic = getattr(self.opt, 'lambda_dssim', 0.2) * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + target_tensor = self.gaussians.get_anchor if hasattr(self.gaussians, 'get_anchor') else self.gaussians._xyz + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = target_tensor.grad.clone() if target_tensor.grad is not None else torch.zeros_like(target_tensor) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = target_tensor.grad.clone() if target_tensor.grad is not None else torch.zeros_like(target_tensor) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < getattr(self.opt, 'densify_until_iter', 15000): + if hasattr(self.gaussians, 'add_densification_stats'): + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > getattr(self.opt, 'densify_from_iter', 500) and step % getattr(self.opt, 'densification_interval', 100) == 0: + if hasattr(self.gaussians, 'densify_and_prune'): + self.gaussians.densify_and_prune(getattr(self.opt, 'densify_grad_threshold', 0.0002), 0.005, self.scene.cameras_extent, 20) + + if step % getattr(self.opt, 'opacity_reset_interval', 3000) == 0: + if hasattr(self.gaussians, 'reset_opacity'): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_anchors = target_tensor.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_anchors), "delta_N": int(num_anchors - self.last_n_anchors), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_anchors = num_anchors + + histograms = {} + if step % 1000 == 0: + if hasattr(self.gaussians, 'get_opacity'): + histograms["opacity"] = torch.sigmoid(self.gaussians.get_opacity).clone().detach() + if hasattr(self.gaussians, 'get_scaling'): + scales = torch.exp(self.gaussians.get_scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians.get_anchor if hasattr(self.gaussians, 'get_anchor') else self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + if hasattr(self.gaussians, 'get_scaling'): + raw_scales = self.gaussians.get_scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + + if hasattr(self.gaussians, 'get_opacity'): + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians.get_opacity))) + + if cameras is not None and len(cameras) > 0 and hasattr(self.gaussians, 'get_rotation'): + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians.get_rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + + if not hasattr(self.gaussians, 'get_anchor') or not hasattr(self.gaussians, 'get_scaling'): + return densities + + xyz = self.gaussians.get_anchor + opacities = torch.sigmoid(self.gaussians.get_opacity).squeeze() + scales = torch.exp(self.gaussians.get_scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + diff --git a/methods/wrapper_sgfgs.py b/methods/wrapper_sgfgs.py new file mode 100644 index 0000000000000000000000000000000000000000..ef41667bcb10e31e249f6b510234bf0028d09f58 --- /dev/null +++ b/methods/wrapper_sgfgs.py @@ -0,0 +1,271 @@ +import os +import sys +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../sgfgs_official'))) +from utils.loss_utils import l1_loss, ssim, predicted_normal_loss, delta_normal_loss, zero_one_loss +from utils.sgf_utils import SGFHelper, guided_l1_loss, guided_ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("sgfgs") +class SGFGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree, self.dataset.brdf_dim, self.dataset.brdf_mode, self.dataset.brdf_envmap_res) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.sgf_helper = SGFHelper() + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + if self.pipe.brdf: + self.gaussians.set_requires_grad("normal", state=step >= self.opt.normal_reg_from_iter) + self.gaussians.set_requires_grad("normal2", state=step >= self.opt.normal_reg_from_iter) + if self.gaussians.brdf_mode == "envmap": + self.gaussians.brdf_mlp.build_mips() + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + losses_extra = {} + if self.pipe.brdf and step > self.opt.normal_reg_from_iter: + if step < self.opt.normal_reg_util_iter: + losses_extra['predicted_normal'] = predicted_normal_loss(render_pkg["normal"], render_pkg["normal_ref"], render_pkg["alpha"]) + losses_extra['zero_one'] = zero_one_loss(render_pkg["alpha"]) + if "delta_normal_norm" in render_pkg.keys(): + losses_extra['delta_reg'] = delta_normal_loss(render_pkg["delta_normal_norm"], render_pkg["alpha"]) + + loss_target = torch.tensor(0.0, device="cuda") + loss_parasitic = torch.tensor(0.0, device="cuda") + struct_discrepancy = self.sgf_helper.compute_structure_discrepancy(image, gt_image) + + if self.opt.use_sgf: + gamma = getattr(self.opt, 'feedback_nonlinear_gamma', 1.2) + beta = getattr(self.opt, 'feedback_p_weighting_beta', 0.5) + guidance_map = self.sgf_helper.get_guidance_map(struct_discrepancy, gamma=gamma, beta=beta) + + alpha_l1 = getattr(self.opt, 'alpha_l1_feedback', 0.7) + alpha_ssim = getattr(self.opt, 'alpha_ssim_feedback', 1.0) + lambda_struct = getattr(self.opt, 'lambda_struct_loss', 0.05) + + w_l1 = alpha_l1 * guidance_map + w_ssim = alpha_ssim * guidance_map + + Ll1 = guided_l1_loss(image, gt_image, w_l1) + loss_ssim_guided = guided_ssim(image, gt_image, w_ssim) + + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * loss_ssim_guided + loss_struct = struct_discrepancy.mean() + loss_parasitic = lambda_struct * loss_struct + else: + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + for k in losses_extra.keys(): + loss_parasitic += getattr(self.opt, f'lambda_{k}') * losses_extra[k] + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0 and loss_parasitic > 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + if self.pipe.brdf and self.pipe.brdf_mode == "envmap": + self.gaussians.brdf_mlp.clamp_(min=0.0, max=1.0) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_sogs.py b/methods/wrapper_sogs.py new file mode 100644 index 0000000000000000000000000000000000000000..97ed0a13f0fd436733a380d6418860f520aa67d4 --- /dev/null +++ b/methods/wrapper_sogs.py @@ -0,0 +1,338 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +_UPSTREAM = '/root/autodl-tmp/SOGS_offy' +sys.path.insert(0, _UPSTREAM) +# monorepo-flat: insert each first-level subdir individually +import os as _os +for _sub in ('utils', 'scene', 'gaussian_renderer', 'arguments', 'lpipsPyTorch', 'experiments'): + _p = _os.path.join(_UPSTREAM, _sub) + if _os.path.isdir(_p): + sys.path.insert(0, _p) +del _os, _p, _sub, _UPSTREAM +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel + +class ConfigNode(dict): + def __init__(self, **kwargs): + super().__init__() + for k, v in kwargs.items(): + if isinstance(v, dict): + self[k] = ConfigNode(**v) + else: + self[k] = v + def __getattr__(self, key): + try: + return self[key] + except KeyError: + raise AttributeError(f"'ConfigNode' object has no attribute '{key}'") + def __setattr__(self, key, value): + self[key] = value + +@register_method("sogs") +class SOGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.args = ConfigNode( + source_path=dataset_config["source_path"], + model_path=dataset_config["model_path"], + eval=True, + resolution=dataset_config.get("resolution", 1), + images="images", + data_device="cuda", + white_background=False, + sh_degree=3 + ) + + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.opt = ConfigNode( + iterations=30000, + position_lr_init=0.00016, + position_lr_final=0.0000016, + position_lr_delay_mult=0.01, + position_lr_max_steps=30000, + feature_lr=0.0025, + opacity_lr=0.05, + scaling_lr=0.005, + rotation_lr=0.001, + percent_dense=0.01, + lambda_dssim=0.2, + densification_interval=100, + opacity_reset_interval=3000, + densify_from_iter=500, + densify_until_iter=15000, + densify_grad_threshold=0.0002, + random_background=False, + disable_xyz_log_activation=False + ) + + self.pipe = ConfigNode( + convert_SHs_python=False, + compute_cov3D_python=False, + debug=False + ) + + self.sorting = ConfigNode( + enabled=True, + weights=ConfigNode(xyz=1.0), + improvement_break=0.001, + shuffle=False, + normalize=True, + activated=True + ) + + self.neighbor_loss = ConfigNode( + lambda_neighbor=1.0, + weights=ConfigNode(opacity=1.0), + blur=ConfigNode(kernel_size=5, sigma=1.0), + loss_fn="mse", + normalize=False, + activated=True + ) + + self.run = ConfigNode(no_progress_bar=True, use_sh=True) + + self.gaussians = GaussianModel(self.args.sh_degree, disable_xyz_log_activation=self.opt.disable_xyz_log_activation) + self.scene = Scene(self.args, self.gaussians) + self.gaussians.training_setup(self.opt) + + if self.sorting.enabled: + self.gaussians.prune_to_square_shape() + self.gaussians.sort_into_grid(self.sorting, not self.run.no_progress_bar) + + bg_color = [1, 1, 1] if self.args.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0 and self.run.use_sh: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + + if self.neighbor_loss.lambda_neighbor > 0: + nb_losses = [] + attr_getter_fn = self.gaussians.get_activated_attr_flat if self.neighbor_loss.activated else self.gaussians.get_attr_flat + weight_sum = sum(self.neighbor_loss.weights.values()) + for attr_name, attr_weight in self.neighbor_loss.weights.items(): + if attr_weight > 0: + nb_losses.append(self.gaussians.neighborloss_2d(attr_getter_fn(attr_name), self.neighbor_loss) * attr_weight / weight_sum) + loss_parasitic = self.neighbor_loss.lambda_neighbor * sum(nb_losses) + else: + loss_parasitic = torch.tensor(0.0, device="cuda") + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0 and loss_parasitic.item() > 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v_t = state.get("exp_avg_sq", torch.zeros_like(p)) + lr = [pg["lr"] for pg in self.gaussians.optimizer.param_groups if pg["params"][0] is p][0] + u_t = (lr / (torch.sqrt(v_t) + 1e-15)) * p.grad.clone() + group_grads.append(u_t.reshape(-1)) + grads_target[group_name] = torch.cat(group_grads) if group_grads else None + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v_t = state.get("exp_avg_sq", torch.zeros_like(p)) + lr = [pg["lr"] for pg in self.gaussians.optimizer.param_groups if pg["params"][0] is p][0] + u_p = (lr / (torch.sqrt(v_t) + 1e-15)) * p.grad.clone() + group_grads.append(u_p.reshape(-1)) + grads_parasitic[group_name] = torch.cat(group_grads) if group_grads else None + + for group_name in param_groups_map: + gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if self.sorting.enabled: + self.gaussians.prune_to_square_shape() + self.gaussians.sort_into_grid(self.sorting, not self.run.no_progress_bar) + + if step % self.opt.opacity_reset_interval == 0 or (self.args.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "nb_loss": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + if hasattr(self.gaussians, "grid_sidelen"): + metrics["grid_sidelen"] = int(self.gaussians.grid_sidelen) + + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + bg = torch.tensor([0,0,0], dtype=torch.float32, device="cuda") + render_pkg = native_render(camera, self.gaussians, self.pipe, bg) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + diff --git a/methods/wrapper_sortfreegs.py b/methods/wrapper_sortfreegs.py new file mode 100644 index 0000000000000000000000000000000000000000..d82fb4bb284000e579eadf86830ee067bd175548 --- /dev/null +++ b/methods/wrapper_sortfreegs.py @@ -0,0 +1,279 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../sort-free-gs_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("sortfreegs") +class SortFreeGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + # Fix: depth_correct not populated by ModelParams.extract + self.dataset.depth_correct = False + + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + loss_target = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target_weighted = (1.0 - self.opt.lambda_dssim) * loss_target + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target_weighted + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target_weighted.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity, self.gaussians._features_do, self.gaussians._features_ro], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target_weighted.backward(retain_graph=True) + grads_target = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get("exp_avg_sq", torch.zeros_like(p)) + lr = self._get_lr_for_param(p) + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + group_grads.append(u.reshape(-1)) + if group_grads: + grads_target[group_name] = torch.cat(group_grads) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + group_grads = [] + for p in params: + if p.grad is not None: + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get("exp_avg_sq", torch.zeros_like(p)) + lr = self._get_lr_for_param(p) + u = (lr / (torch.sqrt(v) + 1e-8)) * p.grad.clone() + group_grads.append(u.reshape(-1)) + if group_grads: + grads_parasitic[group_name] = torch.cat(group_grads) + + for group_name in param_groups_map: + gt = grads_target.get(group_name) + gp = grads_parasitic.get(group_name) + if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sigma_val": float(self.gaussians._sigma.mean()), + "weight_bg_val": float(self.gaussians._weight_background.mean()) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def _get_lr_for_param(self, target_param): + for group in self.gaussians.optimizer.param_groups: + for p in group['params']: + if p is target_param: + return group['lr'] + return 0.0 + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_specgaussian.py b/methods/wrapper_specgaussian.py new file mode 100644 index 0000000000000000000000000000000000000000..831ada88b1b7123edbf3ba02a078d2d5d5d382ed --- /dev/null +++ b/methods/wrapper_specgaussian.py @@ -0,0 +1,318 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/SpecGaussian_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render, prefilter_voxel +from scene import Scene, GaussianModel, SpecularModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("specgaussian") +class SpecGaussianWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.dataset.is_real = "Synthetic" not in self.dataset.source_path + self.use_filter = self.dataset.is_real + + self.gaussians = GaussianModel(self.dataset.sh_degree, self.dataset.asg_degree) + self.specular_mlp = SpecularModel(self.dataset.is_real, self.dataset.is_indoor) + self.specular_mlp.train_setting(self.opt) + + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_effective_step(self, params, grads): + u_list = [] + for p, g in zip(params, grads): + if g is None: + continue + v_t = torch.ones_like(p) + lr = 1e-4 + state = self.gaussians.optimizer.state.get(p, None) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + for group in self.gaussians.optimizer.param_groups: + if any(p is param for param in group["params"]): + lr = group["lr"] + break + u = (lr / (torch.sqrt(v_t) + 1e-8)) * g + u_list.append(u.reshape(-1)) + if len(u_list) == 0: + return torch.zeros(1, device="cuda") + return torch.cat(u_list) + + def train_iteration(self, step): + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + voxel_visible_mask = None + if self.use_filter: + voxel_visible_mask = prefilter_voxel(viewpoint_cam, self.gaussians, self.pipe, self.background) + + mlp_color = 0.0 + if step > 3000: + dir_pp = (self.gaussians.get_xyz - viewpoint_cam.camera_center.repeat(self.gaussians.get_features.shape[0], 1)) + dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True) + normal = self.gaussians.get_normal_axis(dir_pp_normalized=dir_pp_normalized, return_delta=True) + if self.use_filter and voxel_visible_mask is not None: + mlp_color = self.specular_mlp.step(self.gaussians.get_asg_features[voxel_visible_mask], + dir_pp_normalized[voxel_visible_mask], + normal.detach()[voxel_visible_mask]) + else: + mlp_color = self.specular_mlp.step(self.gaussians.get_asg_features, dir_pp_normalized, normal.detach()) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, mlp_color, voxel_visible_mask=voxel_visible_mask) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest, self.gaussians._features_asg], + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.specular_mlp.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + grads_target_raw = {group_name: [p.grad.clone() if p.grad is not None else torch.zeros_like(p) for p in params] for group_name, params in param_groups_map.items()} + u_target = {group_name: self._get_effective_step(params, grads_target_raw[group_name]) for group_name, params in param_groups_map.items()} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.specular_mlp.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + grads_parasitic_raw = {group_name: [p.grad.clone() if p.grad is not None else torch.zeros_like(p) for p in params] for group_name, params in param_groups_map.items()} + u_parasitic = {group_name: self._get_effective_step(params, grads_parasitic_raw[group_name]) for group_name, params in param_groups_map.items()} + + u_target_global = torch.cat(list(u_target.values())) + u_parasitic_global = torch.cat(list(u_parasitic.values())) + + valid_mask = (torch.abs(u_target_global) > 0) | (torch.abs(u_parasitic_global) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_target_global[valid_mask].unsqueeze(0), u_parasitic_global[valid_mask].unsqueeze(0))[0]) + parasitic_ratio = float(u_parasitic_global.norm() / (u_target_global.norm() + u_parasitic_global.norm() + 1e-7)) + + for group_name in param_groups_map: + ut = u_target[group_name] + up = u_parasitic[group_name] + if ut.norm() > 0 and up.norm() > 0: + cos = float(F.cosine_similarity(ut.unsqueeze(0), up.unsqueeze(0))[0]) + r = float(up.norm() / (ut.norm() + up.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.specular_mlp.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if self.use_filter: + self.gaussians.max_radii2D[voxel_visible_mask] = torch.max(self.gaussians.max_radii2D[voxel_visible_mask], radii[visibility_filter]) + else: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + + if step < self.opt.densify_until_iter: + viewspace_point_tensor_densify = render_pkg["viewspace_points_densify"] + self.gaussians.add_densification_stats(viewspace_point_tensor_densify, visibility_filter, voxel_visible_mask, self.use_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.update_learning_rate(step) + self.specular_mlp.optimizer.step() + self.specular_mlp.optimizer.zero_grad(set_to_none=True) + self.specular_mlp.update_learning_rate(step) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "specular_intensity": float(mlp_color.norm(dim=-1).mean()) if isinstance(mlp_color, torch.Tensor) else 0.0 + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + histograms["asg_mag"] = self.gaussians._features_asg.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + voxel_visible_mask = None + if self.use_filter: + voxel_visible_mask = prefilter_voxel(camera, self.gaussians, self.pipe, self.background) + + dir_pp = (self.gaussians.get_xyz - camera.camera_center.repeat(self.gaussians.get_features.shape[0], 1)) + dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True) + normal = self.gaussians.get_normal_axis(dir_pp_normalized=dir_pp_normalized, return_delta=True) + + mlp_color = 0.0 + if self.use_filter and voxel_visible_mask is not None: + mlp_color = self.specular_mlp.step(self.gaussians.get_asg_features[voxel_visible_mask], + dir_pp_normalized[voxel_visible_mask], normal[voxel_visible_mask]) + else: + mlp_color = self.specular_mlp.step(self.gaussians.get_asg_features, dir_pp_normalized, normal) + + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, mlp_color, voxel_visible_mask=voxel_visible_mask) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + self.specular_mlp.save_weights(save_dir, step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc = self.gaussians._features_dc + rest = self.gaussians._features_rest + asg = self.gaussians._features_asg + energy_rest = rest.norm(dim=-1).mean() if rest is not None and rest.shape[1] > 0 else 0.0 + energy_asg = asg.norm(dim=-1).mean() if asg is not None and asg.shape[1] > 0 else 0.0 + metrics["sh_energy_ratio"] = float((energy_rest + energy_asg) / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_splatfactow.py b/methods/wrapper_splatfactow.py new file mode 100644 index 0000000000000000000000000000000000000000..a2ee0047030f69d17c083bb41aea4e3548891f97 --- /dev/null +++ b/methods/wrapper_splatfactow.py @@ -0,0 +1,298 @@ +import os +import sys +import torch +import numpy as np +import torch.nn.functional as F +from pathlib import Path +from plyfile import PlyData, PlyElement + +from core.registry import register_method +from core.base_method import BaseMethod + +from splatfactow.splatfactow_config import splatfactow_config +from nerfstudio.data.datamanagers.full_images_datamanager import FullImageDatamanagerConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig + +@register_method("splatfactow") +class SplatfactowWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.config = splatfactow_config.config + res = dataset_config.get("resolution", 1) + + self.config.pipeline.datamanager = FullImageDatamanagerConfig( + dataparser=NerfstudioDataParserConfig( + load_3D_points=True, + data=Path(dataset_config["source_path"]), + downscale_factor=res, + eval_mode="fraction" + ), + cache_images_type="uint8" + ) + self.config.pipeline.datamanager.data = Path(dataset_config["source_path"]) + + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.pipeline = self.config.pipeline.setup(device="cuda", test_mode="train") + self.pipeline.datamanager.setup_train() + self.optimizers = self.config.optimizers.setup(pipeline=self.pipeline) + + self.model = self.pipeline.model + self.last_n_gaussians = self.model.num_points + + def train_iteration(self, step): + camera, data = self.pipeline.datamanager.next_train(step) + + outputs = self.model(camera) + loss_dict = self.model.get_loss_dict(outputs, data) + + loss_target = loss_dict.get("main_loss", torch.tensor(0.0, device="cuda")) + loss_parasitic = loss_dict.get("scale_reg", torch.tensor(0.0, device="cuda")) + \ + loss_dict.get("alpha_loss", torch.tensor(0.0, device="cuda")) + \ + loss_dict.get("entropy_loss", torch.tensor(0.0, device="cuda")) + + loss = loss_target + loss_parasitic + if "camera_opt_regularization" in loss_dict: + loss += loss_dict["camera_opt_regularization"] + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + self.optimizers.zero_grad_all() + loss_target.backward(retain_graph=True) + + grad_target_global = [] + for param_group in self.model.get_param_groups().values(): + for p in param_group: + if p.grad is not None: + grad_target_global.append(p.grad.clone().reshape(-1)) + if grad_target_global: + grad_target_global = torch.cat(grad_target_global) + else: + grad_target_global = torch.zeros(1, device="cuda") + + self.optimizers.zero_grad_all() + loss_parasitic.backward(retain_graph=True) + + grad_parasitic_global = [] + for param_group in self.model.get_param_groups().values(): + for p in param_group: + if p.grad is not None: + grad_parasitic_global.append(p.grad.clone().reshape(-1)) + if grad_parasitic_global: + grad_parasitic_global = torch.cat(grad_parasitic_global) + else: + grad_parasitic_global = torch.zeros(1, device="cuda") + + if grad_target_global.norm() > 0 and grad_parasitic_global.norm() > 0: + grad_cos_sim = float(F.cosine_similarity(grad_target_global.unsqueeze(0), grad_parasitic_global.unsqueeze(0))) + parasitic_ratio = float(grad_parasitic_global.norm() / (grad_target_global.norm() + 1e-7)) + + param_groups_map = { + "spatial": ["means"], + "geometry": ["scales", "quats"], + "opacity": ["opacities"], + "appearance": ["appearance_features", "colors"] + } + + self.optimizers.zero_grad_all() + loss_target.backward(retain_graph=True) + grads_t = {k: [] for k in param_groups_map} + for group_name, keys in param_groups_map.items(): + for key in keys: + if key in self.model.gauss_params and self.model.gauss_params[key].grad is not None: + opt = self.optimizers.optimizers.get(key) + if opt is not None: + p = self.model.gauss_params[key] + lr = opt.param_groups[0]["lr"] + state = opt.state[p] + vt = state.get("exp_avg_sq", torch.zeros_like(p)) + u = (lr / (torch.sqrt(vt) + 1e-8)) * p.grad.clone() + grads_t[group_name].append(u.reshape(-1)) + + self.optimizers.zero_grad_all() + loss_parasitic.backward(retain_graph=True) + grads_p = {k: [] for k in param_groups_map} + for group_name, keys in param_groups_map.items(): + for key in keys: + if key in self.model.gauss_params and self.model.gauss_params[key].grad is not None: + opt = self.optimizers.optimizers.get(key) + if opt is not None: + p = self.model.gauss_params[key] + lr = opt.param_groups[0]["lr"] + state = opt.state[p] + vt = state.get("exp_avg_sq", torch.zeros_like(p)) + u = (lr / (torch.sqrt(vt) + 1e-8)) * p.grad.clone() + grads_p[group_name].append(u.reshape(-1)) + + for group_name in param_groups_map: + gt_list = grads_t[group_name] + gp_list = grads_p[group_name] + if gt_list and gp_list: + gt = torch.cat(gt_list) + gp = torch.cat(gp_list) + if gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.optimizers.zero_grad_all() + loss.backward() + else: + self.optimizers.zero_grad_all() + loss.backward() + + self.optimizers.optimizer_step("all") + self.optimizers.scheduler_step("all") + + self.model.step = step + self.model.after_train(step) + self.model.refinement_after(self.optimizers, step) + + num_gaussians = self.model.num_points + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_dict.get("ssim", 0.0)), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.model.opacities).clone().detach() + scales = torch.exp(self.model.scales).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + if "colors" in self.model.gauss_params: + histograms["sh_dc_mag"] = self.model.gauss_params["colors"].detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + outputs = self.model.get_outputs_for_camera(camera) + return {"image": outputs["rgb"], "depth": outputs.get("depth", None)} + + def save(self, save_dir, step): + os.makedirs(save_dir, exist_ok=True) + filename = os.path.join(save_dir, "point_cloud.ply") + + count = self.model.num_points + xyz = self.model.means.detach().cpu().numpy() + normals = np.zeros_like(xyz) + + f_dc = self.model.shs_0.contiguous().detach().cpu().numpy() + f_rest = self.model.shs_rest.transpose(1, 2).contiguous().detach().cpu().reshape(count, -1).numpy() + opacities = self.model.opacities.detach().cpu().numpy() + scales = self.model.scales.detach().cpu().numpy() + rots = self.model.quats.detach().cpu().numpy() + + dtype_full = [('x', 'f4'), ('y', 'f4'), ('z', 'f4'), ('nx', 'f4'), ('ny', 'f4'), ('nz', 'f4')] + for i in range(f_dc.shape[1]): + dtype_full.append((f'f_dc_{i}', 'f4')) + for i in range(f_rest.shape[1]): + dtype_full.append((f'f_rest_{i}', 'f4')) + dtype_full.append(('opacity', 'f4')) + for i in range(scales.shape[1]): + dtype_full.append((f'scale_{i}', 'f4')) + for i in range(rots.shape[1]): + dtype_full.append((f'rot_{i}', 'f4')) + + elements = np.empty(count, dtype=dtype_full) + elements['x'] = xyz[:, 0] + elements['y'] = xyz[:, 1] + elements['z'] = xyz[:, 2] + elements['nx'] = normals[:, 0] + elements['ny'] = normals[:, 1] + elements['nz'] = normals[:, 2] + + for i in range(f_dc.shape[1]): + elements[f'f_dc_{i}'] = f_dc[:, i] + for i in range(f_rest.shape[1]): + elements[f'f_rest_{i}'] = f_rest[:, i] + + elements['opacity'] = opacities[:, 0] + + for i in range(scales.shape[1]): + elements[f'scale_{i}'] = scales[:, i] + for i in range(rots.shape[1]): + elements[f'rot_{i}'] = rots[:, i] + + el = PlyElement.describe(elements, 'vertex') + PlyData([el]).write(filename) + + ckpt_dir = os.path.join(os.path.dirname(os.path.dirname(save_dir)), "nerfstudio_models") + os.makedirs(ckpt_dir, exist_ok=True) + ckpt_path = os.path.join(ckpt_dir, f"step-{step:09d}.ckpt") + torch.save({"step": step, "pipeline": self.pipeline.state_dict()}, ckpt_path) + + def load(self, model_path, iteration): + ckpt_path = os.path.join(model_path, "nerfstudio_models", f"step-{iteration:09d}.ckpt") + if os.path.exists(ckpt_path): + loaded_state = torch.load(ckpt_path, map_location="cpu") + self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"]) + + def get_spatial_centers(self): + return self.model.means.detach() + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.model.scales) + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S = torch.max(scales_2d, dim=1)[0] + min_S = torch.min(scales_2d, dim=1)[0] + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.model.opacities))) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.model.quats.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.model.means + opacities = torch.sigmoid(self.model.opacities).squeeze() + scales = torch.exp(self.model.scales) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_steepgs.py b/methods/wrapper_steepgs.py new file mode 100644 index 0000000000000000000000000000000000000000..0ef2bb4d3ca0d02620fbc119d1bb49e42120811b --- /dev/null +++ b/methods/wrapper_steepgs.py @@ -0,0 +1,258 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../SteepGS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("steepgs") +class SteepGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.densify_strategy = ['steepest'] + self.opt.densify_S_threshold = -1e-6 + self.opt.S_estimator = "partial" + self.opt.xyz_grad_estimator = "mean" + self.opt.xyz_grad_ema_ratio = 0.5 + self.opt.densification_interval = 100 + self.opt.opacity_reset_interval = 3000 + self.opt.densify_from_iter = 500 + self.opt.densify_until_iter = 15000 + self.opt.densify_grad_threshold = 0.0002 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.last_saddle_points = 0 + self.last_split_points = 0 + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + splitting_mats = render_pkg["splitting_matrices"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + loss_target = (1.0 - self.opt.lambda_dssim) * l1_loss(image, gt_image) + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim(image, gt_image)) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz) + if state is not None and "exp_avg_sq" in state: + v_t = state["exp_avg_sq"] + lr = self.gaussians.get_current_learning_rate("xyz") + if lr is not None: + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, splitting_mats, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + densify_info = self.gaussians.densify_and_prune( + self.opt.densify_strategy, + self.opt.densify_grad_threshold, + self.opt.densify_S_threshold, + 0.005, + self.scene.cameras_extent, + size_threshold + ) + if isinstance(densify_info, dict): + self.last_saddle_points = densify_info.get("num_saddle_points", 0) + self.last_split_points = densify_info.get("num_split_points", 0) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "saddle_points": int(self.last_saddle_points), + "split_points": int(self.last_split_points) + } + + if self.track_decoupling: + metrics["grad_cos_sim"] = float(grad_cos_sim) + metrics["parasitic_ratio"] = float(parasitic_ratio) + + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_sugar.py b/methods/wrapper_sugar.py new file mode 100644 index 0000000000000000000000000000000000000000..a998282411cce1e8672e151295aae45ff5410209 --- /dev/null +++ b/methods/wrapper_sugar.py @@ -0,0 +1,342 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +@register_method("sugar") +class SuGaRWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + import sys + conflicting_modules = ['scene', 'arguments', 'gaussian_renderer', 'utils', 'sugar_scene', 'sugar_utils', 'sugar_trainers'] + for mod in list(sys.modules.keys()): + for conflict in conflicting_modules: + if mod == conflict or mod.startswith(conflict + '.'): + sys.modules.pop(mod, None) + + sugar_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../sugar_official')) + gs_path = os.path.join(sugar_path, 'gaussian_splatting') + if gs_path not in sys.path: sys.path.insert(0, gs_path) + if sugar_path not in sys.path: sys.path.insert(0, sugar_path) + + from arguments import ModelParams, PipelineParams, OptimizationParams as VanillaOptParams + from scene import Scene, GaussianModel + from gaussian_renderer import render as vanilla_render + from utils.loss_utils import l1_loss, ssim + from sugar_scene.sugar_model import SuGaR + from sugar_scene.sugar_optimizer import OptimizationParams as SugarOptParams, SuGaROptimizer + from sugar_scene.sugar_densifier import SuGaRDensifier + + self._vanilla_render = vanilla_render + self._l1_loss = l1_loss + self._ssim = ssim + self._SuGaR = SuGaR + self._SugarOptParams = SugarOptParams + self._SuGaROptimizer = SuGaROptimizer + self._SuGaRDensifier = SuGaRDensifier + + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = VanillaOptParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.args.white_background = dataset_config.get("white_background", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(self.bg_color, dtype=torch.float32, device="cuda") + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.phase = 1 + self.sugar = None + self.sugar_optimizer = None + self.sugar_densifier = None + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _transition_to_sugar(self): + from sugar_scene.cameras import CamerasWrapper + self.phase = 2 + points = self.gaussians.get_xyz.detach().float().cuda() + colors = self.gaussians.get_features[:, 0].detach().float().cuda() + + self.sugar = self._SuGaR( + nerfmodel=None, + points=points, + colors=colors, + initialize=False, + sh_levels=self.dataset.sh_degree + 1, + learnable_positions=True, + triangle_scale=1.0, + keep_track_of_knn=True, + knn_to_track=16, + beta_mode='average', + freeze_gaussians=False, + surface_mesh_to_bind=None, + device=torch.device("cuda") + ) + + self.sugar._scales.data[...] = self.gaussians._scaling.detach() + self.sugar._quaternions.data[...] = self.gaussians._rotation.detach() + self.sugar.all_densities.data[...] = self.gaussians._opacity.detach() + self.sugar._sh_coordinates_dc.data[...] = self.gaussians._features_dc.detach() + self.sugar._sh_coordinates_rest.data[...] = self.gaussians._features_rest.detach() + + self.cameras_wrapper = CamerasWrapper(self.scene.getTrainCameras()) + opt_params = self._SugarOptParams( + iterations=15000, position_lr_init=0.00016, position_lr_final=0.0000016, + position_lr_delay_mult=0.01, position_lr_max_steps=30000, feature_lr=0.0025, + opacity_lr=0.05, scaling_lr=0.005, rotation_lr=0.001 + ) + self.sugar_optimizer = self._SuGaROptimizer(self.sugar, opt_params, spatial_lr_scale=5.0) + self.sugar_densifier = self._SuGaRDensifier(self.sugar, self.sugar_optimizer, 0.0002, 0.005, 20, 5.0, 0.01) + self.sugar.train() + + def train_iteration(self, step): + if step == 7001 and self.phase == 1: + self._transition_to_sugar() + + if self.phase == 1: + return self._train_phase_1(step) + else: + return self._train_phase_2(step) + + def _train_phase_1(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = self._vanilla_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + gt_image = viewpoint_cam.original_image.cuda() + loss_l1 = self._l1_loss(image, gt_image) + loss_ssim_val = 1.0 - self._ssim(image, gt_image) + loss = (1.0 - self.opt.lambda_dssim) * loss_l1 + self.opt.lambda_dssim * loss_ssim_val + + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[render_pkg["visibility_filter"]] = torch.max(self.gaussians.max_radii2D[render_pkg["visibility_filter"]], render_pkg["radii"][render_pkg["visibility_filter"]]) + self.gaussians.add_densification_stats(render_pkg["viewspace_points"], render_pkg["visibility_filter"]) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, 20 if step > self.opt.opacity_reset_interval else None) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + current_n = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_l1), "loss_ssim": float(loss_ssim_val), + "num_gaussians": int(current_n), "delta_N": int(current_n - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) + } + self.last_n_gaussians = current_n + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + histograms["anisotropy"] = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def _train_phase_2(self, step): + self.sugar_optimizer.update_learning_rate(step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + cam_idx = self.scene.getTrainCameras().index(viewpoint_cam) + outputs = self.sugar.render_image_gaussian_rasterizer( + nerf_cameras=self.cameras_wrapper, camera_indices=cam_idx, bg_color=self.background, + sh_deg=self.dataset.sh_degree, compute_color_in_rasterizer=True, + compute_covariance_in_rasterizer=True, return_2d_radii=True, return_opacities=True + ) + + # outputs['image']: (H,W,3) -> (3,H,W); original_image is already (3,H,W) + img_out = outputs['image'] + if img_out.shape[-1] == 3 and img_out.dim() == 3: + pred_rgb = img_out.permute(2, 0, 1) + else: + pred_rgb = img_out + gt_rgb = viewpoint_cam.original_image.cuda() + if pred_rgb.shape != gt_rgb.shape: + import torch.nn.functional as F + pred_rgb = F.interpolate(pred_rgb.unsqueeze(0), size=gt_rgb.shape[1:], mode="bilinear", align_corners=False).squeeze(0) + + loss_l1 = self._l1_loss(pred_rgb, gt_rgb) + loss_ssim_val = 1.0 - self._ssim(pred_rgb, gt_rgb) + loss_target = 0.8 * loss_l1 + 0.2 * loss_ssim_val + + loss_parasitic = torch.tensor(0.0, device="cuda") + if step > 7000 and step < 9000: + vis_opacities = outputs['opacities'][outputs['radii'] > 0] if (outputs['radii'] > 0).any() else outputs['opacities'] + loss_parasitic = loss_parasitic + 0.1 * (-(vis_opacities * torch.log(vis_opacities + 1e-10) + (1 - vis_opacities) * torch.log(1 - vis_opacities + 1e-10))).mean() + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.sugar_optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.sugar._points.grad.clone() if self.sugar._points.grad is not None else torch.zeros_like(self.sugar._points) + + self.sugar_optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.sugar._points.grad.clone() if self.sugar._points.grad is not None else torch.zeros_like(self.sugar._points) + + state = self.sugar_optimizer.optimizer.state.get(self.sugar._points, None) + if state is not None and 'exp_avg_sq' in state: + v_t = state['exp_avg_sq'] + eta = self.sugar_optimizer.optimizer.param_groups[0]['lr'] + u_t = (eta / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (eta / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.sugar_optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + self.sugar_optimizer.step() + self.sugar_optimizer.zero_grad(set_to_none=True) + + current_n = self.sugar.n_points + metrics = { + "loss": float(loss), "loss_l1": float(loss_l1), "loss_ssim": float(loss_ssim_val), + "num_gaussians": int(current_n), "delta_N": int(current_n - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "loss_entropy": float(loss_parasitic) + } + self.last_n_gaussians = current_n + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.sugar.all_densities).clone().detach() + scales = torch.exp(self.sugar._scales).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + histograms["anisotropy"] = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["sh_dc_mag"] = self.sugar._sh_coordinates_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + if self.phase == 1: + render_pkg = self._vanilla_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"]} + else: + cam_idx = self.scene.getTrainCameras().index(camera) + outputs = self.sugar.render_image_gaussian_rasterizer( + nerf_cameras=self.cameras_wrapper, camera_indices=cam_idx, bg_color=self.background, + sh_deg=self.dataset.sh_degree, compute_color_in_rasterizer=True + ) + return {"image": outputs["image"].transpose(0, 1).transpose(1, 2)} + + def save(self, save_dir, step): + if self.phase == 1: + self.scene.save(step) + else: + model_path = os.path.join(save_dir, f"sugar_{step}.pt") + self.sugar.save_model(path=model_path) + + def load(self, model_path, iteration): + pass + + def get_spatial_centers(self): + if self.phase == 1: + return self.gaussians._xyz + return self.sugar.points + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + if self.phase == 1: + scales = torch.exp(self.gaussians._scaling) + opacities = torch.sigmoid(self.gaussians._opacity) + else: + scales = torch.exp(self.sugar._scales) + opacities = torch.sigmoid(self.sugar.all_densities) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(opacities)) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + if self.phase == 1: + return torch.zeros(query_points.shape[0], device="cuda") + + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + chunk_size = 50000 + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + densities[i:end] = self.sugar.compute_density(query_points[i:end]) + return densities diff --git a/methods/wrapper_sundae.py b/methods/wrapper_sundae.py new file mode 100644 index 0000000000000000000000000000000000000000..44793aa43d9018921f2e403c3b924425b314f6d9 --- /dev/null +++ b/methods/wrapper_sundae.py @@ -0,0 +1,267 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../SUNDAE_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("sundae") +class SundaeWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + sti_metrics = {} + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + + steps_target = {} + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + for group_name, params in param_groups_map.items(): + st_list = [] + for p in params: + if p.grad is not None: + g = p.grad.clone() + state = self.gaussians.optimizer.state.get(p, None) + v = state.get("exp_avg_sq", torch.ones_like(g)) if state else torch.ones_like(g) + step_val = g / (torch.sqrt(v) + 1e-8) + st_list.append(step_val.reshape(-1)) + if st_list: + steps_target[group_name] = torch.cat(st_list) + + steps_parasitic = {} + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + for group_name, params in param_groups_map.items(): + sp_list = [] + for p in params: + if p.grad is not None: + g = p.grad.clone() + state = self.gaussians.optimizer.state.get(p, None) + v = state.get("exp_avg_sq", torch.ones_like(g)) if state else torch.ones_like(g) + step_val = g / (torch.sqrt(v) + 1e-8) + sp_list.append(step_val.reshape(-1)) + if sp_list: + steps_parasitic[group_name] = torch.cat(sp_list) + + for group_name in param_groups_map: + st = steps_target.get(group_name) + sp = steps_parasitic.get(group_name) + if st is not None and sp is not None and st.norm() > 0 and sp.norm() > 0: + cos = float(F.cosine_similarity(st.unsqueeze(0), sp.unsqueeze(0))) + r = float(sp.norm() / (st.norm() + sp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + sti_metrics[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(sti_metrics) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, "point_cloud", f"iteration_{iteration}", "point_cloud.ply")) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_tamgs.py b/methods/wrapper_tamgs.py new file mode 100644 index 0000000000000000000000000000000000000000..81cabfda47b977f2f1118bc121e239abe8e57857 --- /dev/null +++ b/methods/wrapper_tamgs.py @@ -0,0 +1,399 @@ +import os +import sys +import torch +import random +from random import randint +import torch.nn.functional as F + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../taming-3dgs'))) + +from utils.loss_utils import l1_loss +from fused_ssim import fused_ssim as fast_ssim +sys.path.insert(0, '/root/autodl-tmp/taming-3dgs') + +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from argparse import ArgumentParser +from utils.taming_utils import compute_gaussian_score, get_edges, get_count_array + +@register_method("tamgs") +class TamGSWrapper(BaseMethod): + + SEMANTIC_GROUPS = { + "spatial": [("xyz", "_xyz")], + "geometry": [("scaling", "_scaling"), ("rotation", "_rotation")], + "opacity": [("opacity", "_opacity")], + "appearance": [("f_dc", "_features_dc"), ("f_rest", "_features_rest")], + } + + BUDGET_MULTIPLIERS = { + "bicycle": 15, "flowers": 15, "garden": 15, "stump": 15, "treehill": 15, + "room": 2, "counter": 2, "kitchen": 2, "bonsai": 2, + "truck": 2, "train": 2, + "playroom": 5, "drjohnson": 5, + } + + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.parser.add_argument("--cams", type=int, default=10) + self.parser.add_argument("--budget", type=float, default=2.0) + self.parser.add_argument("--mode", type=str, default="multiplier") + self.parser.add_argument("--ho_iteration", type=int, default=15000) + self.parser.add_argument("--sh_lower", action='store_true', default=False) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + + if "Synthetic_NeRF" in self.args.source_path or "blender" in self.args.source_path.lower(): + self.args.white_background = True + + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + scene_key = os.path.basename(self.args.source_path.rstrip("/")).lower() + self.args.budget = float(self.BUDGET_MULTIPLIERS.get(scene_key, 2)) + self.args.densification_interval = 500 + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.opt.densification_interval = self.args.densification_interval + + self.gaussians = GaussianModel(self.dataset.sh_degree, self.opt.optimizer_type) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.score_coefficients = { + 'view_importance': 50, 'edge_importance': 50, 'mse_importance': 50, + 'grad_importance': 25, 'dist_importance': 50, 'opac_importance': 100, + 'dept_importance': 5, 'loss_importance': 10, 'radii_importance': 10, + 'scale_importance': 25, 'count_importance': 0.1, 'blend_importance': 50 + } + + self.all_edges = [] + for view in self.scene.getTrainCameras(): + edges_loss = get_edges(view.original_image).squeeze().cuda() + edges_loss_norm = (edges_loss - torch.min(edges_loss)) / (torch.max(edges_loss) - torch.min(edges_loss)) + self.all_edges.append(edges_loss_norm.cpu()) + + self.counts_array = get_count_array( + len(self.gaussians.get_xyz), self.args.budget, self.opt, mode=self.args.mode) + self.densify_iter_num = 0 + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.ho_triggered = False + + def _get_optimizer_for_group(self, opt_group_name): + if opt_group_name == "f_rest": + return self.gaussians.shoptimizer + return self.gaussians.optimizer + + def _extract_adam_state(self, opt, group_name): + for pg in opt.param_groups: + if pg["name"] == group_name: + param = pg["params"][0] + lr = pg["lr"] + state = opt.state.get(param, {}) + v = state.get("exp_avg_sq", None) + return lr, v + return 0.0, None + + def _effective_update(self, grad, lr, v, eps=1e-15): + if v is None or lr == 0.0: + return torch.zeros_like(grad) + return (lr / (torch.sqrt(v) + eps)) * grad + + def _compute_sti_probes(self, loss_target, loss_parasitic, loss_total): + results = {} + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if self.gaussians.shoptimizer: + self.gaussians.shoptimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + target_grads = {} + for gname, params in self.SEMANTIC_GROUPS.items(): + for oname, aname in params: + p = getattr(self.gaussians, aname) + target_grads[(gname, oname)] = p.grad.clone() if p.grad is not None else torch.zeros_like(p) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if self.gaussians.shoptimizer: + self.gaussians.shoptimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + parasitic_grads = {} + for gname, params in self.SEMANTIC_GROUPS.items(): + for oname, aname in params: + p = getattr(self.gaussians, aname) + parasitic_grads[(gname, oname)] = p.grad.clone() if p.grad is not None else torch.zeros_like(p) + + for gname, params in self.SEMANTIC_GROUPS.items(): + u_t_parts, u_p_parts = [], [] + for oname, aname in params: + opt = self._get_optimizer_for_group(oname) + lr, v = self._extract_adam_state(opt, oname) + g_t = target_grads[(gname, oname)] + g_p = parasitic_grads[(gname, oname)] + u_t_parts.append(self._effective_update(g_t, lr, v).flatten()) + u_p_parts.append(self._effective_update(g_p, lr, v).flatten()) + + u_t_vec = torch.cat(u_t_parts) + u_p_vec = torch.cat(u_p_parts) + + active = u_p_vec.abs() > 1e-10 + if active.sum() < 2: + results[f"sti_{gname}_cos"] = 0.0 + results[f"sti_{gname}_ratio"] = 0.0 + results[f"sti_{gname}_ti"] = 0.0 + continue + + m_t = u_t_vec[active] + m_p = u_p_vec[active] + s = F.cosine_similarity(m_t.unsqueeze(0), m_p.unsqueeze(0)).item() + norm_mt = m_t.norm() + norm_mp = m_p.norm() + r_prime = float(norm_mp / (norm_mt + norm_mp + 1e-10)) + ti = r_prime * max(0.0, -s) + + results[f"sti_{gname}_cos"] = s + results[f"sti_{gname}_ratio"] = r_prime + results[f"sti_{gname}_ti"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + if self.gaussians.shoptimizer: + self.gaussians.shoptimizer.zero_grad(set_to_none=True) + loss_total.backward() + + return results + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(randint(0, len(self.viewpoint_stack) - 1)) + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = fast_ssim(image.unsqueeze(0), gt_image.unsqueeze(0)) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + + sti_metrics = {} + if self.track_decoupling and step % 100 == 0 and step > 0: + sti_metrics = self._compute_sti_probes(loss_target, loss_parasitic, loss) + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.iterations: + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max( + self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + target_total = self.counts_array[self.densify_iter_num + 1] + current_N = self.gaussians.get_xyz.shape[0] + + if current_N >= target_total - 10: + self.densify_iter_num += 1 + else: + train_cams = self.scene.getTrainCameras() + edge_losses, camlist = [], [] + num_cams = self.args.cams if self.args.cams != -1 else 10 + for _ in range(num_cams): + loc = random.randint(0, len(train_cams) - 1) + target_cam = train_cams[loc] + camlist.append(target_cam) + edge_map = self.all_edges[loc].cuda() + h, w = target_cam.image_height, target_cam.image_width + if edge_map.shape[0] != h or edge_map.shape[1] != w: + edge_map = F.interpolate( + edge_map.unsqueeze(0).unsqueeze(0), + size=(h, w), mode='bilinear', align_corners=False + ).squeeze() + edge_losses.append(edge_map) + + gaussian_importance = compute_gaussian_score( + self.scene, camlist, edge_losses, self.gaussians, + self.pipe, self.background, self.score_coefficients, self.opt) + + safe_budget = int(max(target_total, current_N + 50)) + try: + self.gaussians.densify_with_score( + scores=gaussian_importance, max_screen_size=20, + min_opacity=0.005, extent=self.scene.cameras_extent, + budget=safe_budget, radii=radii, + iter_num=self.densify_iter_num) + except Exception: + pass + self.densify_iter_num += 1 + + if step % self.opt.opacity_reset_interval == 0 or \ + (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + if step == self.args.ho_iteration and not self.ho_triggered: + self.gaussians.modify_functions() + self.ho_triggered = True + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.shoptimizer.step() + self.gaussians.shoptimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + delta_N = num_gaussians - self.last_n_gaussians + self.last_n_gaussians = num_gaussians + + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": num_gaussians, + "delta_N": delta_N, + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + } + + if self.track_decoupling: + for k, v in sti_metrics.items(): + metrics[k] = float(v) + + budget_idx = min(self.densify_iter_num, len(self.counts_array) - 1) + budget_target = int(self.counts_array[budget_idx]) + metrics["budget_target"] = budget_target + metrics["budget_utilization"] = float(num_gaussians) / max(budget_target, 1) + metrics["ho_active"] = 1 if self.ho_triggered else 0 + metrics["densify_iter_num"] = self.densify_iter_num + + histograms = {} + if step % 1000 == 0: + with torch.no_grad(): + histograms["opacity"] = self.gaussians.get_opacity.clone().detach().squeeze() + scales = self.gaussians.get_scaling.clone().detach() + histograms["scaling"] = scales.mean(dim=1) + histograms["anisotropy"] = scales.max(dim=1)[0] / (scales.min(dim=1)[0] + 1e-7) + sh_dc = self.gaussians.get_features_dc.clone().detach() + sh_rest = self.gaussians.get_features_rest.clone().detach() + if sh_rest is not None and sh_rest.shape[1] > 0: + histograms["sh_dc_mag"] = sh_dc.norm(dim=-1).squeeze() + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply( + os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = self.gaussians.get_scaling.clone() + max_s, _ = scales.max(dim=1) + min_s, _ = scales.min(dim=1) + gamma = max_s / (min_s + 1e-7) + metrics["gamma_median"] = float(gamma.median()) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(scales.mean()) + metrics["alpha_mean"] = float(self.gaussians.get_opacity.mean()) + + dc = self.gaussians._features_dc.clone() + rest = self.gaussians._features_rest.clone() + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float( + rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize( + torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([ + 2 * (x * z + w * y), + 2 * (y * z - w * x), + 1 - 2 * (x * x + y * y) + ], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points, cameras=None): + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = self.gaussians.get_opacity.squeeze() + scales = self.gaussians.get_scaling + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_g = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_g + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities \ No newline at end of file diff --git a/methods/wrapper_texture_gs.py b/methods/wrapper_texture_gs.py new file mode 100644 index 0000000000000000000000000000000000000000..4f279ae47649b0abcce5a588c085dbcee32d1a70 --- /dev/null +++ b/methods/wrapper_texture_gs.py @@ -0,0 +1,202 @@ +from functools import partial +import os +import sys +import torch +import numpy as np +import logging +from addict import Dict +import importlib.util + +method_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../Texture-GS')) +for p in list(sys.path): + if 'taming-3dgs' in p or '2dgs_official' in p: + sys.path.remove(p) +if method_root not in sys.path: + sys.path.insert(0, method_root) + + +from core.registry import register_method +from core.base_method import BaseMethod +from models import create_model +from render import create_render_func +from dataset import create_dataset + +dummy_logger = logging.getLogger("TextureGS") +dummy_logger.setLevel(logging.INFO) + +@register_method("texture_gs") +class TextureGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.cfg = Dict({ + "model_cfg": { + "type": "Gaussian3D", "sh_degree": 3, + "geo_emb_dim": 128, + "tex_cfg": {"resolution": 1024, "max_sh_degree": 3}, + "uv_net_cfg": {"emb_dim": 128, "pre_mlp_cfg": {"use_tcnn": True, "n_hidden_layers": 1, "n_neurons": 128}, "mlp_cfg": {"use_tcnn": True, "n_hidden_layers": 2, "n_neurons": 128}}, + "inv_uv_net_cfg": {"emb_dim": 128, "n_sample_points": 2048, "pre_mlp_cfg": {"hash_grid_cfg": {"max_hashmap": 12, "n_features_per_level": 4, "n_levels": 8}, "use_tcnn": True, "n_hidden_layers": 1, "n_neurons": 128}, "mlp_cfg": {"use_tcnn": True, "n_hidden_layers": 2, "n_neurons": 128}} + }, + "dataset_cfg": { + "type": "scene", + "data_root_dir": dataset_config["source_path"], + "image_path": "images", + "eval": True, + "resolution": dataset_config.get("resolution", 1), + "background": [0, 0, 0] + }, + "optim_cfg": { + "uv_net_lr": 0.00002, "inv_uv_net_lr": 0.00002, "uv_net_milestones": [10000, 20000], "uv_net_gamma": 0.5, + "tex_optim_range": [35000, 75000], "tex_lr": 0.0025, + "gaussian_optim_range": [0, 30000], "position_lr_init": 0.0001, "position_lr_final": 0.000001, "position_lr_delay_mult": 0.01, "position_lr_max_steps": 30000, + "feature_lr": 0.0025, "opacity_lr": 0.05, "scaling_lr": 0.005, "rotation_lr": 0.001, "percent_dense": 0.01 + }, + "loss_cfg": { + "lambda_dssim": 0.2, "lambda_no_sh": 2.0, "lambda_alpha": 1.0, "lambda_norm_smooth": 0.5, "lambda_01": 0.001, "lambda_norm": 0.1, + "rgb_range": [0, None], "rgb_no_sh_range": [35000, None], "alpha_range": [0, None], "norm_smooth_range": [2000, 30000], + "opacity_reg_range": [2000, 30000], "norm_range": [2000, 30000], "norm_reg_range": [2000, 30000], "inverse_range": [30000, 35000] + }, + "render_cfg": {"type": "render"}, + "train_cfg": { + "num_iterations": 30000, "densification_interval": 100, "opacity_reset_interval": 3000, + "densify_from_iter": 500, "densify_until_iter": 15000, "densify_grad_threshold": 0.0002, + "min_scale_reset_interval": 1000, "min_scale_reset_from_iter": 4000, + "opacity_prune_interval": 6000, "opacity_prune_theshold": 0.5 + }, + "work_dir": dataset_config["model_path"], "debug": False, "init_from": None, "init_uv_map_from": None + }) + + os.makedirs(self.cfg.work_dir, exist_ok=True) + self.scene = create_dataset(self.cfg.dataset_cfg, dummy_logger, self.cfg.work_dir, False) + + if not hasattr(self.scene, 'train_cameras') or len(self.scene.train_cameras) == 0: + print("[TEXTURE_GS] Warning: 底层静默吞并了相机数据，启动主动强行读取！") + + cam_utils_path = os.path.join(method_root, "utils", "cameras.py") + import sys + if method_root not in sys.path: + sys.path.insert(0, method_root) + from utils.cameras import cameraList_from_camInfos + + class DummyArgs: + resolution = dataset_config.get("resolution", 1) + data_device = "cuda" + white_background = False + + if not hasattr(self.scene, 'train_cameras'): + self.scene.train_cameras = {} + + self.scene.train_cameras[1.0] = cameraList_from_camInfos(self.scene.scene_info.train_cameras, 1.0, DummyArgs()) + + if not hasattr(self.scene, 'test_cameras'): + self.scene.test_cameras = {} + if hasattr(self.scene.scene_info, 'test_cameras') and self.scene.scene_info.test_cameras: + self.scene.test_cameras[1.0] = cameraList_from_camInfos(self.scene.scene_info.test_cameras, 1.0, DummyArgs()) + else: + self.scene.test_cameras[1.0] = [] + + def safe_get_train(scale=1.0): + if scale in self.scene.train_cameras: return self.scene.train_cameras[scale] + return list(self.scene.train_cameras.values())[0] + self.scene.getTrainCameras = safe_get_train + + def safe_get_test(scale=1.0): + if hasattr(self.scene, 'test_cameras') and scale in self.scene.test_cameras: return self.scene.test_cameras[scale] + if hasattr(self.scene, 'test_cameras') and self.scene.test_cameras: return list(self.scene.test_cameras.values())[0] + return [] + self.scene.getTestCameras = safe_get_test + + self.gaussians = create_model(self.cfg.model_cfg, dummy_logger, self.cfg.work_dir, False) + self.render_func = create_render_func(self.cfg.render_cfg) + self.gaussians.initialize(self.scene.scene_info.point_cloud, self.scene.cameras_extent) + self.gaussians.setup_optim(self.cfg.optim_cfg) + + self.background = torch.tensor([0, 0, 0], dtype=torch.float32, device="cuda") + # render_func binding + self.render_func = partial(self.render_func, cfg=self.cfg.render_cfg, bg_color=self.background) + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.current_stage = 1 + + def transition_to_stage2(self): + print("\n[TEXTURE_GS] >>> 自动换挡: 阶段 2 (UV 映射 30k-35k) <<<") + tmp_path = os.path.join(self.cfg.work_dir, "stage1_tmp.pth") + try: + state_dict = self.gaussians.capture() + except AttributeError: + state_dict = (self.gaussians._xyz, self.gaussians._features_dc, self.gaussians._features_rest, self.gaussians._scaling, self.gaussians._rotation, self.gaussians._opacity, self.gaussians.max_radii2D, self.gaussians.xyz_gradient_accum, self.gaussians.denom, self.gaussians.optimizer.state_dict(), self.gaussians.spatial_lr_scale) + torch.save(({'params': state_dict}, 30000), tmp_path) + + self.cfg.init_from = tmp_path + self.cfg.model_cfg.type = "UVMapGaussian3D" + self.gaussians = create_model(self.cfg.model_cfg, dummy_logger, self.cfg.work_dir, False) + self.gaussians.initialize(self.scene.scene_info.point_cloud, self.scene.cameras_extent) + self.gaussians.setup_optim(self.cfg.optim_cfg) + self.current_stage = 2 + + def transition_to_stage3(self): + print("\n[TEXTURE_GS] >>> 自动换挡: 阶段 3 (纹理生成 35k-75k) <<<") + tmp_path_uv = os.path.join(self.cfg.work_dir, "stage2_tmp.pth") + uv_dict = {} + if hasattr(self.gaussians, 'uv_net'): uv_dict['uv_net'] = self.gaussians.uv_net.state_dict() + if hasattr(self.gaussians, 'inv_uv_net'): uv_dict['inv_uv_net'] = self.gaussians.inv_uv_net.state_dict() + torch.save((uv_dict, 35000), tmp_path_uv) + + self.cfg.init_uv_map_from = tmp_path_uv + self.cfg.model_cfg.type = "TextureGaussian3D" + self.gaussians = create_model(self.cfg.model_cfg, dummy_logger, self.cfg.work_dir, False) + self.gaussians.initialize(self.scene.scene_info.point_cloud, self.scene.cameras_extent) + self.gaussians.setup_optim(self.cfg.optim_cfg) + self.cfg.render_cfg.type = "uv_tex_render" + self.render_func = partial(create_render_func(self.cfg.render_cfg), cfg=self.cfg.render_cfg, bg_color=self.background) + self.current_stage = 3 + + def train_iteration(self, step): + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint = self.viewpoint_stack.pop(np.random.randint(0, len(self.viewpoint_stack))) + + if step == 30001 and self.current_stage == 1: + self.transition_to_stage2() + elif step == 35001 and self.current_stage == 2: + self.transition_to_stage3() + + loss_val = 0.0 + if self.current_stage == 1: + loss, _, extra_info = self.gaussians.compute_loss(step, 30000, viewpoint, self.render_func, self.cfg.loss_cfg) + loss_val = float(loss.item()) + if step > 2000 and step % 1000 == 0: self.gaussians.reset_min_scale() + self.gaussians.optimize_step(step, 30000, self.cfg.train_cfg, extra_info) + elif self.current_stage == 2: + loss, _, extra_info = self.gaussians.compute_loss(step, 75000, viewpoint, self.render_func, self.cfg.loss_cfg) + loss_val = float(loss.item()) + self.gaussians.optimizer_uv.step() + self.gaussians.optimizer_uv.zero_grad(set_to_none=True) + elif self.current_stage == 3: + loss, _, extra_info = self.gaussians.compute_loss(step, 75000, viewpoint, self.render_func, self.cfg.loss_cfg) + loss_val = float(loss.item()) + self.gaussians.optimizer_tex.step() + if step > 45000: self.gaussians.optimizer.step() + self.gaussians.optimizer_tex.zero_grad(set_to_none=True) + self.gaussians.optimizer.zero_grad(set_to_none=True) + + return {"loss": loss_val, "num_gaussians": int(len(self.gaussians.get_xyz)), "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3))}, {} + + def render(self, camera): + with torch.no_grad(): + try: + render_pkg = self.render_func(camera, self.gaussians) + img = render_pkg.get("render", render_pkg.get("image")) + except Exception: + img = torch.zeros((3, 256, 256), device="cuda") + return {"image": img} + + def get_spatial_centers(self): return self.gaussians.get_xyz + + def compute_physical_metrics(self, cameras=None): + with torch.no_grad(): + scales = torch.exp(self.gaussians._scaling) + return {"gamma_median": float(torch.median(scales.max(dim=-1)[0]/(scales.min(dim=-1)[0]+1e-7))), "alpha_mean": float(torch.mean(torch.sigmoid(self.gaussians._opacity)))} + + def evaluate_spatial_field(self, query_points, cameras=None): return torch.zeros(query_points.shape[0], device="cuda") + + def save(self, save_dir, step): self.gaussians.save_point_cloud(os.path.join(save_dir, f"iteration_{step}.ply")) + + def load(self, model_path, iteration): pass diff --git a/methods/wrapper_texturedgs.py b/methods/wrapper_texturedgs.py new file mode 100644 index 0000000000000000000000000000000000000000..fe70d836b4a2e9c07caed24bbfdf8c19ca19cbe8 --- /dev/null +++ b/methods/wrapper_texturedgs.py @@ -0,0 +1,283 @@ +import os +import sys +import torch +import torch.nn.functional as F + +PROJ_ROOT = "/root/autodl-tmp/textured_gaussians" +EXAMPLES_ROOT = "/root/autodl-tmp/textured_gaussians/examples" + +if PROJ_ROOT not in sys.path: + sys.path.insert(0, PROJ_ROOT) +if EXAMPLES_ROOT not in sys.path: + sys.path.insert(0, EXAMPLES_ROOT) + +from core.registry import register_method +from core.base_method import BaseMethod +from torchmetrics.image import StructuralSimilarityIndexMeasure + +class MockCamera: + def __init__(self, c2w): + self.camera_center = c2w[:3, 3] + self.world_view_transform = torch.linalg.inv(c2w).transpose(0, 1) + +class MockScene: + def __init__(self, train_set, device): + self.cams = [] + for i in range(len(train_set)): + cam = train_set[i] + if isinstance(cam, dict) and "camtoworld" in cam: + c2w = cam["camtoworld"].squeeze() + self.cams.append(MockCamera(c2w.to(device))) + elif hasattr(cam, "world_view_transform"): + self.cams.append(cam) + else: + self.cams.append(MockCamera(torch.eye(4, device=device))) + + def getTrainCameras(self): + return self.cams + + def getTestCameras(self): + return getattr(self, 'test_cams', []) + +@register_method("texturedgs") +class TexturedGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + from textured_gaussians.rendering import rasterization_textured_gaussians + + self.device = "cuda" + self.source_path = dataset_config["source_path"] + self.model_path = dataset_config["model_path"] + self.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.render_fn = rasterization_textured_gaussians + + # 初始化真实的可导 SSIM 算子 + self.ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0).to(self.device) + + if os.path.exists(os.path.join(self.source_path, "transforms_train.json")): + self.train_set = self._load_blender_robust(self.source_path, self.resolution) + self.scene_scale = 1.0 + else: + from datasets.colmap import Parser, Dataset + self.parser = Parser( + data_dir=self.source_path, + factor=self.resolution, + test_every=8 + ) + self.train_set = Dataset(self.parser, split="train") + self.test_set = Dataset(self.parser, split="test") + self.scene_scale = self.parser.scene_scale + + self.scene = MockScene(self.train_set, self.device) + if hasattr(self, "test_set"): + self.scene.test_cams = MockScene(self.test_set, self.device).cams + + num_pts = 10000 + self.params = torch.nn.ParameterDict({ + "means": torch.nn.Parameter(torch.rand(num_pts, 3, device=self.device) * 2 - 1), + "scales": torch.nn.Parameter(torch.zeros(num_pts, 3, device=self.device)), + "quats": torch.nn.Parameter(torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).repeat(num_pts, 1)), + "opacities": torch.nn.Parameter(torch.zeros(num_pts, 1, device=self.device)), + "sh0": torch.nn.Parameter(torch.zeros(num_pts, 1, 3, device=self.device)), + "textures": torch.nn.Parameter(torch.ones(num_pts, 50, 50, 4, device=self.device)) + }) + + self.optimizer = torch.optim.Adam([{"params": list(self.params.values()), "lr": 0.001}]) + self.last_n_gaussians = num_pts + + def _load_blender_robust(self, data_dir, resolution): + import json, glob, cv2 + import imageio.v2 as imageio + import numpy as np + + with open(os.path.join(data_dir, "transforms_train.json"), "r") as f: + meta = json.load(f) + + camera_angle_x = float(meta["camera_angle_x"]) + cams = [] + + for frame in meta["frames"]: + fpath = os.path.join(data_dir, frame["file_path"]) + if not fpath.endswith(".png"): fpath += ".png" + if not os.path.exists(fpath): + base = os.path.basename(frame["file_path"]) + matches = glob.glob(os.path.join(data_dir, "**", base + "*"), recursive=True) + if matches: fpath = matches[0] + + rgba = imageio.imread(fpath) + H, W = rgba.shape[:2] + if resolution != 1: + H, W = H // resolution, W // resolution + rgba = cv2.resize(rgba, (W, H)) + + image = (torch.tensor(rgba) / 255.0).float() + image = image[..., :3] * image[..., 3:] + (1.0 - image[..., 3:]) + + c2w = torch.tensor(frame["transform_matrix"]).float() + focal = .5 * W / np.tan(.5 * camera_angle_x) + K = torch.tensor([[focal, 0, W/2], [0, focal, H/2], [0, 0, 1]]).float() + + cams.append({ + "image": image.to(self.device), + "camtoworld": c2w.unsqueeze(0).to(self.device), + "K": K.unsqueeze(0).to(self.device), + "width": W, "height": H + }) + return cams + + def train_iteration(self, step): + cam = self.train_set[step % len(self.train_set)] + + if isinstance(cam, dict): + camtoworlds = cam["camtoworld"].unsqueeze(0).to(self.device) + Ks = cam["K"].unsqueeze(0).to(self.device) + image_gt = cam["image"].to(self.device) + if image_gt.max() > 2.0: image_gt = image_gt / 255.0 + height, width = cam["image"].shape[0], cam["image"].shape[1] + else: + camtoworlds = cam["camtoworld"].unsqueeze(0).to(self.device).unsqueeze(0).to(self.device) + Ks = cam["K"].unsqueeze(0).to(self.device).unsqueeze(0).to(self.device) + image_gt = cam["image"].to(self.device) + if image_gt.max() > 2.0: image_gt = image_gt / 255.0.to(self.device) + if image_gt.max() > 2.0: image_gt = image_gt / 255.0 + width, height = image_gt.shape[1], image_gt.shape[0] + + means = self.params["means"] + quats = self.params["quats"] + scales = torch.exp(self.params["scales"]) + opacities = torch.sigmoid(self.params["opacities"]).squeeze(-1) + + textures = self.params["textures"] + rgb_textures = textures[..., :3] + alpha_textures = textures[..., 3:4] + alpha_textures = alpha_textures / (alpha_textures.amax(dim=[1, 2], keepdim=True) + 1e-6) + textures = torch.cat([rgb_textures, alpha_textures], dim=-1).clamp(0.0, 1.0) + + viewmats = torch.linalg.inv(camtoworlds) + + out = self.render_fn( + means=means, quats=quats, scales=scales, opacities=opacities, + colors=self.params["sh0"], textures=textures, + viewmats=viewmats, Ks=Ks, width=width, height=height, + packed=False, sh_degree=0 + ) + image = out[0][0] + + # --- 接入真实的 SSIM 计算 --- + loss_l1 = F.l1_loss(image, image_gt) + + img_bchw = image.unsqueeze(0).permute(0, 3, 1, 2) + gt_bchw = image_gt.unsqueeze(0).permute(0, 3, 1, 2) + loss_ssim = 1.0 - self.ssim_metric(img_bchw, gt_bchw) + + loss = loss_l1 * 0.8 + loss_ssim * 0.2 + # ---------------------------- + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.optimizer.zero_grad(set_to_none=True) + loss_l1.backward(retain_graph=True) + grad_target = self.params["means"].grad.clone() if self.params["means"].grad is not None else torch.zeros_like(self.params["means"]) + + self.optimizer.zero_grad(set_to_none=True) + loss_ssim.backward(retain_graph=True) + grad_parasitic = self.params["means"].grad.clone() if self.params["means"].grad is not None else torch.zeros_like(self.params["means"]) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + self.optimizer.step() + self.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.params["means"].shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_l1), + "loss_ssim": float(loss_ssim), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.params["opacities"]).clone().detach() + scales = torch.exp(self.params["scales"]).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + histograms["anisotropy"] = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["sh_dc_mag"] = self.params["sh0"].detach().norm(dim=-1) + + return metrics, histograms + + def get_spatial_centers(self): + return self.params["means"] + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.params["scales"]) + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_s, _ = torch.max(scales_2d, dim=1) + min_s, _ = torch.min(scales_2d, dim=1) + gamma = max_s / (min_s + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.params["opacities"]))) + + if self.params["textures"].shape[-1] == 4: + metrics["texture_alpha_activation"] = float(torch.mean(self.params["textures"][..., 3])) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + if isinstance(c, dict): view_dirs.append(c.get("camera_center", torch.zeros(3)).tolist()) + else: view_dirs.append(c.camera_center.tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device=self.device), dim=1) + + rots = F.normalize(self.params["quats"].clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + v_points = query_points.shape[0] + densities = torch.zeros(v_points, device=self.device) + xyz = self.params["means"] + opacities = torch.sigmoid(self.params["opacities"]).squeeze() + scales = torch.exp(self.params["scales"]) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + n_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (n_gaussians + 1)) + for i in range(0, v_points, chunk_size): + end = min(i + chunk_size, v_points) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities + + def render(self, camera): + return {"image": torch.zeros(100, 100, 3)} + + def save(self, save_dir, step): pass + def load(self, model_path, iteration): pass diff --git a/methods/wrapper_topologygs.py b/methods/wrapper_topologygs.py new file mode 100644 index 0000000000000000000000000000000000000000..241a1f52e57a6e7cdce8ecd15aa52855cc14bf79 --- /dev/null +++ b/methods/wrapper_topologygs.py @@ -0,0 +1,283 @@ +import os +import sys +import random +import torch +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +_UPSTREAM = '/root/autodl-tmp/Topology-GS_offy' +sys.path.insert(0, _UPSTREAM) +# Some upstream repos are monorepo-flat: scene/, utils/, gaussian_renderer/ +# are siblings under repo root. Nested `from utils.X import ...` needs +# each first-level dir as its own sys.path entry. +import os as _os +for _sub in ('utils', 'scene', 'gaussian_renderer', 'arguments', 'lpipsPyTorch', 'experiments'): + _p = _os.path.join(_UPSTREAM, _sub) + if _os.path.isdir(_p): + sys.path.insert(0, _p) +del _os, _p, _sub, _UPSTREAM +from experiments.config import cfg +from scene import Scene, GaussianModel +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import prefilter_voxel, render as native_render +try: + from TOPO_users.persistence_loss import PersistLoss +except ImportError: + PersistLoss = None + +@register_method("topologygs") +class TopologyGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.dataset_config = dataset_config + self.track_decoupling = hyperparams.get("track_decoupling", False) + + cfg.MODEL.SOURCE_PATH = dataset_config["source_path"] + cfg.MODEL.MODEL_PATH = dataset_config["model_path"] + self.resolution = dataset_config.get("resolution", 1) + + self.net = cfg.MODEL + self.opt = cfg.OPTIM + self.pipe = cfg.PIPE + self.data = cfg.DATA + + self.gaussians = GaussianModel( + self.net.FEAT_DIM, self.net.N_OFFSETS, self.net.VOXEL_SIZE, + self.net.UPDATE_DEPTH, self.net.UPDATE_INIT_FACTOR, + self.net.UPDATE_HIERACHY_FACTOR, self.net.USE_FEAT_BANK, + self.net.APPEARANCE_DIM, self.net.RATIO, self.net.ADD_OPACITY_DIST, + self.net.ADD_COV_DIST, self.net.ADD_COLOR_DIST + ) + self.scene = Scene(self.opt, self.data, self.net, self.gaussians, shuffle=False) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.net.WHITE_BACKGROUND else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.ps = PersistLoss(dims=self.opt.PERSISTENCE.DIMS, ks=self.opt.PERSISTENCE.KS).cuda() if (self.opt.USE_PERS_LOSS and PersistLoss is not None) else None + self.last_n_gaussians = self.gaussians.get_anchor.shape[0] + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack)-1)) + + voxel_visible_mask = prefilter_voxel(viewpoint_cam, self.gaussians, self.pipe, self.background) + retain_grad = (step < self.opt.UPDATE_UNTIL and step >= 0) + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, visible_mask=voxel_visible_mask, retain_grad=retain_grad) + + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + offset_selection_mask = render_pkg["selection_mask"] + opacity = render_pkg["neural_opacity"] + scaling = render_pkg["scaling"] + + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + ssim_loss = (1.0 - ssim(image, gt_image)) + scaling_reg = scaling.prod(dim=1).mean() + + persist_loss = torch.tensor(0.0, device="cuda") + if self.ps is not None: + if step >= self.opt.PERSISTENCE.START and step <= self.opt.PERSISTENCE.END and step % self.opt.PERSISTENCE.INTERVAL == 0: + image_gt = F.interpolate(gt_image.unsqueeze(0), scale_factor=self.opt.PERSISTENCE.RATIO, mode=self.opt.PERSISTENCE.MODE).squeeze(0) + image_pred = F.interpolate(image.unsqueeze(0), scale_factor=self.opt.PERSISTENCE.RATIO, mode=self.opt.PERSISTENCE.MODE).squeeze(0) + try: + p_loss, _ = self.ps(image_pred, image_gt, viewpoint_cam.persistence) + persist_loss = p_loss + except: + pass + + loss_target = (1.0 - self.opt.LAMBDA_DSSIM) * Ll1 + loss_parasitic = self.opt.LAMBDA_DSSIM * ssim_loss + 0.01 * scaling_reg + if persist_loss > 0: + loss_parasitic += self.opt.PERSISTENCE.LAMBDA * persist_loss + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + + if self.track_decoupling and step % 100 == 0: + param_groups_map = { + "spatial": [self.gaussians._anchor, self.gaussians._offset], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._anchor_feat] + } + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + + grads_target = {} + states_v = {} + lrs = {} + + for group_name, params in param_groups_map.items(): + grads_target[group_name] = [] + states_v[group_name] = [] + lrs[group_name] = [] + for p in params: + if p.grad is not None: + grads_target[group_name].append(p.grad.clone().reshape(-1)) + state = self.gaussians.optimizer.state.get(p, {}) + v = state.get('exp_avg_sq', torch.zeros_like(p)).reshape(-1) + states_v[group_name].append(v) + lr_val = 0.0 + for pg in self.gaussians.optimizer.param_groups: + if id(p) in [id(param) for param in pg['params']]: + lr_val = pg['lr'] + break + lrs[group_name].append(torch.full_like(v, lr_val)) + + grads_target[group_name] = torch.cat(grads_target[group_name]) if grads_target[group_name] else torch.empty(0, device="cuda") + states_v[group_name] = torch.cat(states_v[group_name]) if states_v[group_name] else torch.empty(0, device="cuda") + lrs[group_name] = torch.cat(lrs[group_name]) if lrs[group_name] else torch.empty(0, device="cuda") + + grad_target_global = grads_target.get("spatial", torch.empty(0, device="cuda")) + if grad_target_global.numel() == 0: + grad_target_global = torch.zeros(1, device="cuda") + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + + grads_parasitic = {} + for group_name, params in param_groups_map.items(): + grads_parasitic[group_name] = [] + for p in params: + if p.grad is not None: + grads_parasitic[group_name].append(p.grad.clone().reshape(-1)) + grads_parasitic[group_name] = torch.cat(grads_parasitic[group_name]) if grads_parasitic[group_name] else torch.empty(0, device="cuda") + + grad_parasitic_global = grads_parasitic.get("spatial", torch.empty(0, device="cuda")) + if grad_parasitic_global.numel() == 0: + grad_parasitic_global = torch.zeros(1, device="cuda") + + if grad_target_global.norm() > 0 and grad_parasitic_global.norm() > 0: + grad_cos_sim = float(F.cosine_similarity(grad_target_global.unsqueeze(0), grad_parasitic_global.unsqueeze(0))) + parasitic_ratio = float(grad_parasitic_global.norm() / (grad_target_global.norm() + 1e-7)) + + for group_name in param_groups_map: + gt = grads_target.get(group_name, torch.empty(0, device="cuda")) + gp = grads_parasitic.get(group_name, torch.empty(0, device="cuda")) + v = states_v.get(group_name, torch.empty(0, device="cuda")) + lr = lrs.get(group_name, torch.empty(0, device="cuda")) + + if gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + eff = lr / (torch.sqrt(v) + 1e-8) + u_t = eff * gt + u_p = eff * gp + cos = float(F.cosine_similarity(u_t.unsqueeze(0), u_p.unsqueeze(0))) + r = float(u_p.norm() / (u_t.norm() + u_p.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.UPDATE_UNTIL and step > self.opt.START_STAT: + self.gaussians.training_statis(viewspace_point_tensor, opacity, visibility_filter, offset_selection_mask, voxel_visible_mask) + if step > self.opt.UPDATE_FROM and step % self.opt.UPDATE_INTERVAL == 0: + self.gaussians.adjust_anchor(check_interval=self.opt.UPDATE_INTERVAL, success_threshold=self.opt.SUCCESS_THRESHOLD, grad_threshold=self.opt.DENSIFY_GRAD_THRESHOLD, min_opacity=self.opt.MIN_OPACITY) + elif step == self.opt.UPDATE_UNTIL: + torch.cuda.empty_cache() + + if step < self.opt.ITERATIONS: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_anchor.shape[0] * self.net.N_OFFSETS + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), + "parasitic_ratio": float(parasitic_ratio), + "persist_loss": float(persist_loss) + } + metrics.update(stats) + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._anchor_feat.detach().norm(dim=-1) + + return metrics, histograms + + def load(self, model_path, iteration): + (model_params, _) = torch.load(os.path.join(model_path, f"chkpnt{iteration}.pth")) + self.gaussians.restore(model_params, self.opt) + + def get_spatial_centers(self): + return self.gaussians.get_anchor + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians.get_anchor + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_trimgs.py b/methods/wrapper_trimgs.py new file mode 100644 index 0000000000000000000000000000000000000000..4616aa48aa84b73523277fd5095324118c4e9e53 --- /dev/null +++ b/methods/wrapper_trimgs.py @@ -0,0 +1,303 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../TrimGS_official/Trim2DGS'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, TrimGSOptimizationParams + +@register_method("trimgs") +class TrimGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = TrimGSOptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + if "360" in str(self.args.source_path): + self.opt.lambda_dist = 100.0 + else: + self.opt.lambda_dist = 1000.0 + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def culling(self, xyz, cams, expansion=2): + cam_centers = torch.stack([c.camera_center for c in cams], 0).to(xyz.device) + span_x = cam_centers[:, 0].max() - cam_centers[:, 0].min() + span_y = cam_centers[:, 1].max() - cam_centers[:, 1].min() + span_z = cam_centers[:, 2].max() - cam_centers[:, 2].min() + scene_center = cam_centers.mean(0) + span_x = span_x * expansion + span_y = span_y * expansion + span_z = span_z * expansion + x_min = scene_center[0] - span_x / 2 + x_max = scene_center[0] + span_x / 2 + y_min = scene_center[1] - span_y / 2 + y_max = scene_center[1] + span_y / 2 + z_min = scene_center[2] - span_x / 2 + z_max = scene_center[2] + span_x / 2 + valid_mask = (xyz[:, 0] > x_min) & (xyz[:, 0] < x_max) & \ + (xyz[:, 1] > y_min) & (xyz[:, 1] < y_max) & \ + (xyz[:, 2] > z_min) & (xyz[:, 2] < z_max) + return valid_mask, scene_center + + def prune_low_contribution_gaussians(self, gaussians, cameras, pipe, bg, K=5, prune_ratio=0.1): + top_list = [None] * K + for cam in cameras: + trans = native_render(cam, gaussians, pipe, bg, record_transmittance=True) + if top_list[0] is not None: + m = trans > top_list[0] + if m.any(): + for i in range(K - 1): + top_list[K - 1 - i][m] = top_list[K - 2 - i][m] + top_list[0][m] = trans[m] + else: + top_list = [trans.clone() for _ in range(K)] + contribution = torch.stack(top_list, dim=-1).mean(-1) + tile = torch.quantile(contribution, prune_ratio) + prune_mask = contribution < tile + gaussians.prune_points(prune_mask) + torch.cuda.empty_cache() + + def train_iteration(self, step): + MAX_GAUSSIANS = 2_000_000 + + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + + loss_photometric = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + + lambda_normal = self.opt.lambda_normal if step > 7000 else 0.0 + lambda_dist = self.opt.lambda_dist if step > 3000 else 0.0 + + rend_dist = render_pkg.get("rend_dist", torch.tensor(0.0, device="cuda")) + rend_normal = render_pkg.get("rend_normal", torch.tensor(0.0, device="cuda")) + surf_normal = render_pkg.get("surf_normal", torch.tensor(0.0, device="cuda")) + + if "surf_depth" in render_pkg and self.opt.depth_grad_thresh > 0: + depth_map = render_pkg["surf_depth"][0] + depth_map_for_grad = depth_map[None, None] + sobel_kernel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=depth_map.dtype, device=depth_map.device).view(1, 1, 3, 3) + sobel_kernel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=depth_map.dtype, device=depth_map.device).view(1, 1, 3, 3) + depth_map_for_grad = F.pad(depth_map_for_grad, pad=(1, 1, 1, 1), mode="replicate") + depth_grad_x = F.conv2d(depth_map_for_grad, sobel_kernel_x) / 3 + depth_grad_y = F.conv2d(depth_map_for_grad, sobel_kernel_y) / 3 + depth_grad_mag = torch.sqrt(depth_grad_x ** 2 + depth_grad_y ** 2) + depth_grad_weight = (depth_grad_mag < self.opt.depth_grad_thresh).float() + if self.opt.depth_grad_mask_dilation > 0: + mask_di = self.opt.depth_grad_mask_dilation + depth_grad_weight = -1 * F.max_pool2d(-1 * depth_grad_weight, mask_di * 2 + 1, stride=1, padding=mask_di) + depth_grad_weight = depth_grad_weight.squeeze().detach() + normal_error = (1 - (rend_normal * surf_normal).sum(dim=0)) * depth_grad_weight + normal_error = normal_error[None] + else: + normal_error = (1 - (rend_normal * surf_normal).sum(dim=0))[None] + + normal_loss = lambda_normal * normal_error.mean() + dist_loss = lambda_dist * rend_dist.mean() + + loss_target = loss_photometric + loss_parasitic = normal_loss + dist_loss + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0 and self.gaussians.get_xyz.shape[0] < MAX_GAUSSIANS: + grads = self.gaussians.xyz_gradient_accum / self.gaussians.denom + grads[grads.isnan()] = 0.0 + self.gaussians.densify_and_clone(grads, self.opt.densify_grad_threshold, self.scene.cameras_extent) + self.gaussians.densify_and_split(grads, self.opt.densify_grad_threshold, self.scene.cameras_extent) + scene_mask, _ = self.culling(self.gaussians.get_xyz, self.scene.getTrainCameras()) + if self.gaussians.get_xyz.shape[0] < MAX_GAUSSIANS: + self.gaussians.densify_and_scale_split(self.opt.densify_grad_threshold, self.opt.opacity_cull, self.scene.cameras_extent, self.opt.max_screen_size, self.opt.densify_scale_factor, scene_mask, N=3, no_grad=True) + + if step > self.opt.contribution_prune_from_iter and step % self.opt.contribution_prune_interval == 0: + self.prune_low_contribution_gaussians(self.gaussians, self.scene.getTrainCameras().copy(), self.pipe, self.background, K=5, prune_ratio=self.opt.contribution_prune_ratio) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "loss_normal": float(normal_loss), "loss_dist": float(dist_loss) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("surf_depth", render_pkg.get("depth", None))} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanilla.py b/methods/wrapper_vanilla.py new file mode 100644 index 0000000000000000000000000000000000000000..a0a6a845ef20e004ada485e09e4371aeec1a1ebc --- /dev/null +++ b/methods/wrapper_vanilla.py @@ -0,0 +1,293 @@ +import os +import time +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim + +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +import math +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } +@register_method("vanilla_3dgs") +class Vanilla3DGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, {}) + v_t = state.get('exp_avg_sq', torch.zeros_like(self.gaussians._xyz)) + lr = self.opt.position_lr_init + for param_group in self.gaussians.optimizer.param_groups: + if param_group['name'] == 'xyz': + lr = param_group['lr'] + break + + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "iter_time_ms": float((time.perf_counter() - _iter_start) * 1000), + "vram_allocated_GB": float(torch.cuda.memory_allocated() / 1024**3), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanilla_logged.py b/methods/wrapper_vanilla_logged.py new file mode 100644 index 0000000000000000000000000000000000000000..89e98e6420f88f972e5cfeb97ef60e3c1aee88ae --- /dev/null +++ b/methods/wrapper_vanilla_logged.py @@ -0,0 +1,405 @@ +import os +import time +import json +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim + +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +import math +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + +def _sobel_for_logging(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8).squeeze() + + + +def log_gradient_intensity(model_out, step, viewpoint_cam, gaussians, sobel_map): + """After loss.backward(), bucket position-grad magnitude by edge (top-25%) / flat (bottom-25%).""" + xyz_grad = gaussians._xyz.grad + if xyz_grad is None or sobel_map is None: + return + xyz = gaussians.get_xyz.detach() + proj = viewpoint_cam.full_proj_transform + xyz_h = torch.cat([xyz, torch.ones(xyz.shape[0], 1, device=xyz.device)], dim=-1) + clip = xyz_h @ proj + z = clip[:, 3] + valid_z = z > 1e-4 + ndc = clip[:, :2] / (z.unsqueeze(-1) + 1e-7) + if sobel_map.dim() == 3: + sobel_map = sobel_map.squeeze(0) + H, Wd = sobel_map.shape[-2], sobel_map.shape[-1] + px = ((ndc[:, 0] + 1) * 0.5 * Wd).clamp(0, Wd - 1).long() + py = ((ndc[:, 1] + 1) * 0.5 * H).clamp(0, H - 1).long() + in_view = valid_z & (px >= 0) & (px < Wd) & (py >= 0) & (py < H) + sampled = sobel_map[py, px].detach().cpu() + grad_norm = xyz_grad.detach().norm(dim=-1).cpu() + in_view_cpu = in_view.cpu() + all_W = sampled[in_view_cpu] + if len(all_W) < 10: + return + q75 = float(all_W.quantile(0.75)) + q25 = float(all_W.quantile(0.25)) + edge_mask = in_view_cpu & (sampled >= q75) + flat_mask = in_view_cpu & (sampled <= q25) + edge_mean = float(grad_norm[edge_mask].mean()) if edge_mask.any() else 0.0 + flat_mean = float(grad_norm[flat_mask].mean()) if flat_mask.any() else 0.0 + log_path = os.path.join(model_out, "gradient_intensity.json") + log = [] + if os.path.exists(log_path): + try: + with open(log_path) as f: + log = json.load(f) + except Exception: + log = [] + log.append({"step": int(step), "edge_mean": edge_mean, "flat_mean": flat_mean, + "ratio": (edge_mean / max(flat_mean, 1e-9))}) + with open(log_path, "w") as f: + json.dump(log, f) + + +def log_densification_events(model_out, step, viewpoint_cam, gaussians, sobel_map, will_densify_mask): + xyz = gaussians.get_xyz.detach() + if xyz.shape[0] == 0 or sobel_map is None: + return + proj = viewpoint_cam.full_proj_transform + xyz_h = torch.cat([xyz, torch.ones(xyz.shape[0], 1, device=xyz.device)], dim=-1) + clip = xyz_h @ proj + z = clip[:, 3] + valid_z = z > 1e-4 + ndc = clip[:, :2] / (z.unsqueeze(-1) + 1e-7) + if sobel_map.dim() == 3: + sobel_map = sobel_map.squeeze(0) + H, Wd = sobel_map.shape[-2], sobel_map.shape[-1] + px = ((ndc[:, 0] + 1) * 0.5 * Wd).clamp(0, Wd - 1).long() + py = ((ndc[:, 1] + 1) * 0.5 * H).clamp(0, H - 1).long() + in_view = valid_z & (px >= 0) & (px < Wd) & (py >= 0) & (py < H) + sampled = sobel_map[py, px].detach().cpu() + in_view_cpu = in_view.cpu() + will_cpu = will_densify_mask.cpu() if will_densify_mask is not None else in_view_cpu + all_W = sampled[in_view_cpu] + if len(all_W) < 10: + return + q75 = float(all_W.quantile(0.75)) + q25 = float(all_W.quantile(0.25)) + triggered = sampled[in_view_cpu & will_cpu] + if len(triggered) == 0: + return + n_edge = int((triggered >= q75).sum()) + n_flat = int((triggered <= q25).sum()) + log_path = os.path.join(model_out, "densification_events.json") + log = [] + if os.path.exists(log_path): + try: + with open(log_path) as f: + log = json.load(f) + except Exception: + log = [] + log.append({"step": int(step), "n_will_densify": int(len(triggered)), + "n_edge": n_edge, "n_flat": n_flat, + "edge_thresh": q75, "flat_thresh": q25}) + with open(log_path, "w") as f: + json.dump(log, f) + + +@register_method("vanilla_logged") +class VanillaLoggedWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, {}) + v_t = state.get('exp_avg_sq', torch.zeros_like(self.gaussians._xyz)) + lr = self.opt.position_lr_init + for param_group in self.gaussians.optimizer.param_groups: + if param_group['name'] == 'xyz': + lr = param_group['lr'] + break + + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + if step % 50 == 0: + log_gradient_intensity(self.args.model_path, step, viewpoint_cam, self.gaussians, _sobel_for_logging(gt_image)) + else: + loss.backward() + if step % 50 == 0: + log_gradient_intensity(self.args.model_path, step, viewpoint_cam, self.gaussians, _sobel_for_logging(gt_image)) + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # === Densification event logging (Sobel computed for logging only) === + _sobel = _sobel_for_logging(gt_image) + _grads = self.gaussians.xyz_gradient_accum / (self.gaussians.denom + 1e-7) + _grads[_grads.isnan()] = 0.0 + _grad_norm = _grads.norm(dim=-1).squeeze() + _will = (_grad_norm >= self.opt.densify_grad_threshold) + log_densification_events(self.args.model_path, step, viewpoint_cam, self.gaussians, _sobel, _will) + # === End logging === + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "iter_time_ms": float((time.perf_counter() - _iter_start) * 1000), + "vram_allocated_GB": float(torch.cuda.memory_allocated() / 1024**3), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanilla_thresh.py b/methods/wrapper_vanilla_thresh.py new file mode 100644 index 0000000000000000000000000000000000000000..04bdb26e37b9481c9ae97cbc9c66217aefb662e7 --- /dev/null +++ b/methods/wrapper_vanilla_thresh.py @@ -0,0 +1,301 @@ +import os +import time +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim + +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +import math +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } +@register_method("vanilla_thresh") +class VanillaThreshWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + # ABLATION (count-matching): override densify_grad_threshold via env var + if "DENSIFY_GRAD_THRESHOLD" in os.environ: + try: + self.opt.densify_grad_threshold = float(os.environ["DENSIFY_GRAD_THRESHOLD"]) + print(f"[wrapper_vanilla_thresh] densify_grad_threshold overridden to {self.opt.densify_grad_threshold}") + except ValueError: + pass + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + + state = self.gaussians.optimizer.state.get(self.gaussians._xyz, {}) + v_t = state.get('exp_avg_sq', torch.zeros_like(self.gaussians._xyz)) + lr = self.opt.position_lr_init + for param_group in self.gaussians.optimizer.param_groups: + if param_group['name'] == 'xyz': + lr = param_group['lr'] + break + + u_t = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_target + u_p = (lr / (torch.sqrt(v_t) + 1e-8)) * grad_parasitic + + valid_mask = (torch.norm(u_t, dim=1) > 0) & (torch.norm(u_p, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(u_t[valid_mask], u_p[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(u_p, dim=1).mean() / (torch.norm(u_t, dim=1).mean() + 1e-7)) + + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "iter_time_ms": float((time.perf_counter() - _iter_start) * 1000), + "vram_allocated_GB": float(torch.cuda.memory_allocated() / 1024**3), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + self.last_n_gaussians = num_gaussians + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf.py b/methods/wrapper_vanillasgf.py new file mode 100644 index 0000000000000000000000000000000000000000..71554506bba7b642c3e980b4c1272b56eba6dc77 --- /dev/null +++ b/methods/wrapper_vanillasgf.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf") +class VanillaSGFWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_2scale.py b/methods/wrapper_vanillasgf_2scale.py new file mode 100644 index 0000000000000000000000000000000000000000..053b56d82c0ac4fa6aacdcf27ff251c960f4d70f --- /dev/null +++ b/methods/wrapper_vanillasgf_2scale.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_2scale") +class VanillaSGF2ScaleWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5)]: # 2-scale + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_4scale.py b/methods/wrapper_vanillasgf_4scale.py new file mode 100644 index 0000000000000000000000000000000000000000..841a86957c30dc73fef0df63cbd71c7ff21c71de --- /dev/null +++ b/methods/wrapper_vanillasgf_4scale.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_4scale") +class VanillaSGF4ScaleWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25), (0.125, 0.125)]: # 4-scale + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_hook.py b/methods/wrapper_vanillasgf_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..9cb4a7a19e31755f15760559ef1df7ecf1890378 --- /dev/null +++ b/methods/wrapper_vanillasgf_hook.py @@ -0,0 +1,377 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_hook") +class VanillaSGFHookWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + # ABLATION (hook isolation): + # loss is PLAIN L1+SSIM (no W in forward). + # W only modulates position gradient via backward hook below. + Ll1 = l1_loss(image, gt_image) + ssim_val = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_val) + loss_struct = 0.05 * E_struct.mean() + loss_parasitic = loss_struct + loss = loss_target + loss_parasitic + + # === Per-Gaussian weight: sample W at projected screen-position === + _W_map = W.detach() + if _W_map.dim() == 3: + _W_map = _W_map.squeeze(0) + _H_w, _W_w = _W_map.shape[-2], _W_map.shape[-1] + with torch.no_grad(): + _xyz = self.gaussians._xyz.detach() + _proj = viewpoint_cam.full_proj_transform + _xyz_h = torch.cat([_xyz, torch.ones(_xyz.shape[0], 1, device=_xyz.device)], dim=-1) + _clip = _xyz_h @ _proj + _z = _clip[:, 3:4] + _ndc = _clip[:, :2] / (_z + 1e-7) + _px = ((_ndc[:, 0] + 1) * 0.5 * _W_w).clamp(0, _W_w - 1).long() + _py = ((_ndc[:, 1] + 1) * 0.5 * _H_w).clamp(0, _H_w - 1).long() + _w_per = _W_map[_py, _px] + _in_view = (_z.squeeze() > 1e-4) + _w_per = torch.where(_in_view, _w_per, torch.zeros_like(_w_per)) + _w_per_3 = (1.0 + 0.7 * _w_per).unsqueeze(-1) + + # Register hook on _xyz: scale position grad by per-Gaussian weight; color/opacity untouched + def _hook(grad): + return grad * _w_per_3 + + _hook_handle = self.gaussians._xyz.register_hook(_hook) + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + _hook_handle.remove() + else: + loss.backward() + _hook_handle.remove() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(Ll1), + "sgf_loss_ssim_reweighted": float(1.0 - ssim_val) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_laplacian.py b/methods/wrapper_vanillasgf_laplacian.py new file mode 100644 index 0000000000000000000000000000000000000000..be44b88141129d83b3f3936af49b2201b3c5adb6 --- /dev/null +++ b/methods/wrapper_vanillasgf_laplacian.py @@ -0,0 +1,351 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + # ABLATION: Laplacian operator (function name kept for drop-in compatibility) + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + lap = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + L = F.conv2d(luma, lap, padding=1) + return torch.abs(L) + 1e-8 + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_laplacian") +class VanillaSGFLaplacianWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_log.py b/methods/wrapper_vanillasgf_log.py new file mode 100644 index 0000000000000000000000000000000000000000..c6b22c3fc1d2ddebccb354fe377b36e3e68cefd5 --- /dev/null +++ b/methods/wrapper_vanillasgf_log.py @@ -0,0 +1,357 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + # ABLATION: Laplacian-of-Gaussian (5x5 kernel, name kept for drop-in compatibility) + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + log_kernel = torch.tensor([ + [ 0, 0, -1, 0, 0], + [ 0, -1, -2, -1, 0], + [-1, -2, 16, -2, -1], + [ 0, -1, -2, -1, 0], + [ 0, 0, -1, 0, 0], + ], dtype=torch.float32, device=img.device).view(1, 1, 5, 5) + L = F.conv2d(luma, log_kernel, padding=2) + return torch.abs(L) + 1e-8 + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_log") +class VanillaSGFLoGWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_logged.py b/methods/wrapper_vanillasgf_logged.py new file mode 100644 index 0000000000000000000000000000000000000000..151c133a337d0546c72f08773000c04ad66873fd --- /dev/null +++ b/methods/wrapper_vanillasgf_logged.py @@ -0,0 +1,454 @@ +import os +import json +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + + + +def log_gradient_intensity(model_out, step, viewpoint_cam, gaussians, sobel_map): + """After loss.backward(), bucket position-grad magnitude by edge (top-25%) / flat (bottom-25%).""" + xyz_grad = gaussians._xyz.grad + if xyz_grad is None or sobel_map is None: + return + xyz = gaussians.get_xyz.detach() + proj = viewpoint_cam.full_proj_transform + xyz_h = torch.cat([xyz, torch.ones(xyz.shape[0], 1, device=xyz.device)], dim=-1) + clip = xyz_h @ proj + z = clip[:, 3] + valid_z = z > 1e-4 + ndc = clip[:, :2] / (z.unsqueeze(-1) + 1e-7) + if sobel_map.dim() == 3: + sobel_map = sobel_map.squeeze(0) + H, Wd = sobel_map.shape[-2], sobel_map.shape[-1] + px = ((ndc[:, 0] + 1) * 0.5 * Wd).clamp(0, Wd - 1).long() + py = ((ndc[:, 1] + 1) * 0.5 * H).clamp(0, H - 1).long() + in_view = valid_z & (px >= 0) & (px < Wd) & (py >= 0) & (py < H) + sampled = sobel_map[py, px].detach().cpu() + grad_norm = xyz_grad.detach().norm(dim=-1).cpu() + in_view_cpu = in_view.cpu() + all_W = sampled[in_view_cpu] + if len(all_W) < 10: + return + q75 = float(all_W.quantile(0.75)) + q25 = float(all_W.quantile(0.25)) + edge_mask = in_view_cpu & (sampled >= q75) + flat_mask = in_view_cpu & (sampled <= q25) + edge_mean = float(grad_norm[edge_mask].mean()) if edge_mask.any() else 0.0 + flat_mean = float(grad_norm[flat_mask].mean()) if flat_mask.any() else 0.0 + log_path = os.path.join(model_out, "gradient_intensity.json") + log = [] + if os.path.exists(log_path): + try: + with open(log_path) as f: + log = json.load(f) + except Exception: + log = [] + log.append({"step": int(step), "edge_mean": edge_mean, "flat_mean": flat_mean, + "ratio": (edge_mean / max(flat_mean, 1e-9))}) + with open(log_path, "w") as f: + json.dump(log, f) + + +def log_densification_events(model_out, step, viewpoint_cam, gaussians, sobel_map, will_densify_mask): + """Bucket Gaussians-to-be-densified by edge (top-25%) / flat (bottom-25%) of sobel_map.""" + xyz = gaussians.get_xyz.detach() + if xyz.shape[0] == 0 or sobel_map is None: + return + proj = viewpoint_cam.full_proj_transform + xyz_h = torch.cat([xyz, torch.ones(xyz.shape[0], 1, device=xyz.device)], dim=-1) + clip = xyz_h @ proj + z = clip[:, 3] + valid_z = z > 1e-4 + ndc = clip[:, :2] / (z.unsqueeze(-1) + 1e-7) + if sobel_map.dim() == 3: + sobel_map = sobel_map.squeeze(0) + H, Wd = sobel_map.shape[-2], sobel_map.shape[-1] + px = ((ndc[:, 0] + 1) * 0.5 * Wd).clamp(0, Wd - 1).long() + py = ((ndc[:, 1] + 1) * 0.5 * H).clamp(0, H - 1).long() + in_view = valid_z & (px >= 0) & (px < Wd) & (py >= 0) & (py < H) + sampled = sobel_map[py, px].detach().cpu() + in_view_cpu = in_view.cpu() + will_cpu = will_densify_mask.cpu() if will_densify_mask is not None else in_view_cpu + all_W = sampled[in_view_cpu] + if len(all_W) < 10: + return + q75 = float(all_W.quantile(0.75)) + q25 = float(all_W.quantile(0.25)) + triggered = sampled[in_view_cpu & will_cpu] + if len(triggered) == 0: + return + n_edge = int((triggered >= q75).sum()) + n_flat = int((triggered <= q25).sum()) + log_path = os.path.join(model_out, "densification_events.json") + log = [] + if os.path.exists(log_path): + try: + with open(log_path) as f: + log = json.load(f) + except Exception: + log = [] + log.append({"step": int(step), "n_will_densify": int(len(triggered)), + "n_edge": n_edge, "n_flat": n_flat, + "edge_thresh": q75, "flat_thresh": q25}) + with open(log_path, "w") as f: + json.dump(log, f) + + +@register_method("vanillasgf_logged") +class VanillaSGFLoggedWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + if step % 50 == 0: + log_gradient_intensity(self.args.model_path, step, viewpoint_cam, self.gaussians, W) + else: + loss.backward() + if step % 50 == 0: + log_gradient_intensity(self.args.model_path, step, viewpoint_cam, self.gaussians, W) + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # === Densification event logging === + _grads = self.gaussians.xyz_gradient_accum / (self.gaussians.denom + 1e-7) + _grads[_grads.isnan()] = 0.0 + _grad_norm = _grads.norm(dim=-1).squeeze() + _will = (_grad_norm >= self.opt.densify_grad_threshold) + log_densification_events(self.args.model_path, step, viewpoint_cam, self.gaussians, W, _will) + # === End logging === + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_nolstruct.py b/methods/wrapper_vanillasgf_nolstruct.py new file mode 100644 index 0000000000000000000000000000000000000000..71b9519a3f43e8748eadce0681fc7483c8587126 --- /dev/null +++ b/methods/wrapper_vanillasgf_nolstruct.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_nolstruct") +class VanillaSGFNoLStructWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = torch.tensor(0.0, device="cuda") # ABLATION: no L_struct + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_nomultiscale.py b/methods/wrapper_vanillasgf_nomultiscale.py new file mode 100644 index 0000000000000000000000000000000000000000..dce354f971b1e2c6047cd3e862d5b7c4db992409 --- /dev/null +++ b/methods/wrapper_vanillasgf_nomultiscale.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_nomultiscale") +class VanillaSGFNoMultiscaleWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0)]: # ABLATION: single-scale + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_nosg.py b/methods/wrapper_vanillasgf_nosg.py new file mode 100644 index 0000000000000000000000000000000000000000..546328009733a3f463f712a4cd2ab6b0f9d8d4e5 --- /dev/null +++ b/methods/wrapper_vanillasgf_nosg.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_nosg") +class VanillaSGFNoSGWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p ** 0.5 # ABLATION: no stop-gradient + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_prewitt.py b/methods/wrapper_vanillasgf_prewitt.py new file mode 100644 index 0000000000000000000000000000000000000000..6dd238907d64242d3a1385af347f9a20f45f7ae5 --- /dev/null +++ b/methods/wrapper_vanillasgf_prewitt.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [1, 0, -1], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 1, 1], [0, 0, 0], [-1, -1, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_prewitt") +class VanillaSGFPrewittWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_scharr.py b/methods/wrapper_vanillasgf_scharr.py new file mode 100644 index 0000000000000000000000000000000000000000..06b5603cc74c424b78399fb1dd9cb4b8d77660fd --- /dev/null +++ b/methods/wrapper_vanillasgf_scharr.py @@ -0,0 +1,352 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[3, 0, -3], [10, 0, -10], [3, 0, -3]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[3, 10, 3], [0, 0, 0], [-3, -10, -3]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_scharr") +class VanillaSGFScharrWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # === SGF inline (copied from wrapper_pixelgssgf.py) === + E_struct = 0.0 + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_sens.py b/methods/wrapper_vanillasgf_sens.py new file mode 100644 index 0000000000000000000000000000000000000000..f0ecec7372d19f6238818139c89cb0fb3bc19ee5 --- /dev/null +++ b/methods/wrapper_vanillasgf_sens.py @@ -0,0 +1,250 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, tanfovy=tanfovy, bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, debug=pipe.debug) + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + scales, rotations = None, None + else: + cov3D_precomp = None + scales = pc.get_scaling + rotations = pc.get_rotation + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, means2D=means2D, shs=shs, colors_precomp=None, + opacities=opacity, scales=scales, rotations=rotations, cov3D_precomp=cov3D_precomp) + rendered_image = rendered_image.clamp(0, 1) + return {"render": rendered_image, "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, "radii": radii} + + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1,0,-1],[2,0,-2],[1,0,-1]], dtype=torch.float32, device=img.device).view(1,1,3,3) + wy = torch.tensor([[1,2,1],[0,0,0],[-1,-2,-1]], dtype=torch.float32, device=img.device).view(1,1,3,3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + +def create_window(window_size, channel): + def gaussian(ws, sigma): + gauss = torch.Tensor([math.exp(-(x - ws//2)**2 / float(2*sigma**2)) for x in range(ws)]) + return gauss / gauss.sum() + _1D = gaussian(window_size, 1.5).unsqueeze(1) + _2D = _1D.mm(_1D.t()).float().unsqueeze(0).unsqueeze(0) + return Variable(_2D.expand(channel, 1, window_size, window_size).contiguous()) + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + sigma1_sq = F.conv2d(img1*img1, window, padding=window_size//2, groups=channel) - mu1.pow(2) + sigma2_sq = F.conv2d(img2*img2, window, padding=window_size//2, groups=channel) - mu2.pow(2) + sigma12 = F.conv2d(img1*img2, window, padding=window_size//2, groups=channel) - mu1*mu2 + C1, C2 = 0.01**2, 0.03**2 + return ((2*mu1*mu2+C1)*(2*sigma12+C2)) / ((mu1.pow(2)+mu2.pow(2)+C1)*(sigma1_sq+sigma2_sq+C2)) + + +@register_method("vanillasgf_sens") +class VanillaSGFSensWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + # ── 从环境变量读 SGF 超参数 ── + self.sgf_gamma = float(os.environ.get("SGF_GAMMA", "1.2")) + self.sgf_beta = float(os.environ.get("SGF_BETA", "0.5")) + self.sgf_alpha_l1 = float(os.environ.get("SGF_ALPHA_L1", "0.7")) + self.sgf_alpha_ssim = float(os.environ.get("SGF_ALPHA_SSIM", "1.0")) + self.sgf_lambda_struct = float(os.environ.get("SGF_LAMBDA_STRUCT", "0.05")) + print(f"[SGF-SENS] gamma={self.sgf_gamma} beta={self.sgf_beta} " + f"alpha_l1={self.sgf_alpha_l1} alpha_ssim={self.sgf_alpha_ssim} " + f"lambda_struct={self.sgf_lambda_struct}") + + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + except: pass + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + bg_color = [1,1,1] if self.dataset.white_background else [0,0,0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack)-1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gt_image = viewpoint_cam.original_image.cuda() + + # ── SGF with configurable hyperparameters ── + E_struct = torch.zeros(image.shape[1], image.shape[2], device=image.device) + for w_s, s in [(1.0, 1.0), (0.5, 0.5), (0.25, 0.25)]: + if s == 1.0: + p_s, g_s = image, gt_image + else: + p_s = F.interpolate(image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + g_s = F.interpolate(gt_image.unsqueeze(0), scale_factor=s, mode='bilinear', align_corners=False).squeeze(0) + diff = torch.abs(get_sobel_map(p_s) - get_sobel_map(g_s)) + if s != 1.0: + diff = F.interpolate(diff, size=(image.shape[1], image.shape[2]), mode='bilinear', align_corners=False) + E_struct = E_struct + w_s * diff.squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** self.sgf_gamma # ← γ + W = F_p.detach() ** self.sgf_beta # ← β + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_rw = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + self.sgf_alpha_l1 * W)).mean() # ← α_L1 + + ssim_err = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_rw = (self.opt.lambda_dssim * ssim_err * (1.0 + self.sgf_alpha_ssim * W)).mean() # ← α_SSIM + + loss_struct = self.sgf_lambda_struct * E_struct.mean() # ← λ_struct + + loss = loss_l1_rw + loss_ssim_rw + loss_struct + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_g = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_l1_rw), "loss_ssim": float(loss_ssim_rw), + "num_gaussians": int(num_g), "delta_N": int(num_g - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated()/(1024**3)), + "sgf_w_mean": float(W.mean()), "sgf_loss_struct": float(loss_struct), + "iter_time_ms": float((time.perf_counter()-_iter_start)*1000), + } + self.last_n_gaussians = num_g + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize(); _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), + "render_ms": (time.perf_counter()-_t0)*1000} + + def save(self, save_dir, step): self.scene.save(step) + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + def get_spatial_centers(self): return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + scales = torch.exp(self.gaussians._scaling) + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz = self.gaussians._xyz + opacities = torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + N = xyz.shape[0] + chunk = max(1, 30_000_000 // (N + 1)) + for i in range(0, V, chunk): + end = min(i + chunk, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_wgt.py b/methods/wrapper_vanillasgf_wgt.py new file mode 100644 index 0000000000000000000000000000000000000000..7185e3d224b9dc733f6ab0aee73f599d6ae900fc --- /dev/null +++ b/methods/wrapper_vanillasgf_wgt.py @@ -0,0 +1,342 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_wgt") +class VanillaSGFWGTWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # ABLATION (weight source = static GT edge, only S(I) no S(I_hat)): + E_struct = get_sobel_map(gt_image).squeeze() + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vanillasgf_wraw.py b/methods/wrapper_vanillasgf_wraw.py new file mode 100644 index 0000000000000000000000000000000000000000..dbbbb3632ebf67c5e85ddcffaed738cf31af201c --- /dev/null +++ b/methods/wrapper_vanillasgf_wraw.py @@ -0,0 +1,342 @@ +import os +import sys +import math +import time +import random +import torch +import numpy as np +import torch.nn.functional as F +from torch.autograd import Variable +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../3dgsAtlas_official'))) +from utils.loss_utils import l1_loss, ssim +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams +from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer + + +def native_render(viewpoint_camera, pc, pipe, bg_color, scaling_modifier=1.0): + screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0 + try: + screenspace_points.retain_grad() + except: + pass + + tanfovx = math.tan(viewpoint_camera.FoVx * 0.5) + tanfovy = math.tan(viewpoint_camera.FoVy * 0.5) + + raster_settings = GaussianRasterizationSettings( + image_height=int(viewpoint_camera.image_height), + image_width=int(viewpoint_camera.image_width), + tanfovx=tanfovx, + tanfovy=tanfovy, + bg=bg_color, + scale_modifier=scaling_modifier, + viewmatrix=viewpoint_camera.world_view_transform, + projmatrix=viewpoint_camera.full_proj_transform, + sh_degree=pc.active_sh_degree, + campos=viewpoint_camera.camera_center, + prefiltered=False, + debug=pipe.debug + ) + + rasterizer = GaussianRasterizer(raster_settings=raster_settings) + + means3D = pc.get_xyz + means2D = screenspace_points + opacity = pc.get_opacity + + scales = None + rotations = None + cov3D_precomp = None + + if pipe.compute_cov3D_python: + cov3D_precomp = pc.get_covariance(scaling_modifier) + else: + scales = pc.get_scaling + rotations = pc.get_rotation + + shs = pc.get_features + rendered_image, radii = rasterizer( + means3D=means3D, + means2D=means2D, + shs=shs, + colors_precomp=None, + opacities=opacity, + scales=scales, + rotations=rotations, + cov3D_precomp=cov3D_precomp + ) + + rendered_image = rendered_image.clamp(0, 1) + return { + "render": rendered_image, + "viewspace_points": screenspace_points, + "visibility_filter": radii > 0, + "radii": radii, + } + + +# === SGF helper functions (copied verbatim from wrapper_pixelgssgf.py) === + +def get_sobel_map(img): + luma = (0.299 * img[0] + 0.587 * img[1] + 0.114 * img[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + wy = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32, device=img.device).view(1, 1, 3, 3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8) + + +def create_window(window_size, channel): + def gaussian(window_size, sigma): + gauss = torch.Tensor([math.exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)]) + return gauss / gauss.sum() + _1D_window = gaussian(window_size, 1.5).unsqueeze(1) + _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0) + window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous()) + return window + + +def ssim_map(img1, img2, window_size=11, channel=3): + window = create_window(window_size, channel).to(img1.device) + mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel) + mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel) + mu1_sq = mu1.pow(2) + mu2_sq = mu2.pow(2) + mu1_mu2 = mu1 * mu2 + sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq + sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq + sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2 + C1 = 0.01 ** 2 + C2 = 0.03 ** 2 + map_ssim = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)) + return map_ssim + + +@register_method("vanillasgf_wraw") +class VanillaSGFWRawWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def train_iteration(self, step): + _iter_start = time.perf_counter() + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + + gt_image = viewpoint_cam.original_image.cuda() + + # ABLATION (weight source = raw photometric error, no Sobel filtering): + E_struct = torch.abs(image - gt_image).mean(dim=0) + + F_p = (E_struct / (E_struct.max() + 1e-8)) ** 1.2 + W = F_p.detach() ** 0.5 + + Ll1_map = torch.abs(image - gt_image).mean(dim=0) + loss_l1_reweighted = ((1.0 - self.opt.lambda_dssim) * Ll1_map * (1.0 + 0.7 * W)).mean() + + ssim_error_map = 1.0 - ssim_map(image.unsqueeze(0), gt_image.unsqueeze(0)).squeeze(0).mean(dim=0) + loss_ssim_reweighted = (self.opt.lambda_dssim * ssim_error_map * (1.0 + 1.0 * W)).mean() + + loss_struct = 0.05 * E_struct.mean() + + loss_target = loss_l1_reweighted + loss_parasitic = loss_ssim_reweighted + loss_struct + loss = loss_target + loss_parasitic + + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7)) + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + # vanilla backbone: 2-arg densification stats, no pixels + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + # vanilla backbone: densify_and_prune signature includes radii as last arg + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, radii) + + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio), + "sgf_w_mean": float(W.mean()), + "sgf_e_struct_mean": float(E_struct.mean() if hasattr(E_struct, 'mean') else 0.0), + "sgf_loss_struct": float(loss_struct), + "sgf_loss_l1_reweighted": float(loss_l1_reweighted), + "sgf_loss_ssim_reweighted": float(loss_ssim_reweighted) + } + self.last_n_gaussians = num_gaussians + + metrics["iter_time_ms"] = float((time.perf_counter() - _iter_start) * 1000) + metrics["vram_allocated_GB"] = float(torch.cuda.memory_allocated() / 1024**3) + + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + + return metrics, histograms + + def render(self, camera): + torch.cuda.synchronize() + _t0 = time.perf_counter() + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + torch.cuda.synchronize() + render_ms = (time.perf_counter() - _t0) * 1000 + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None), "render_ms": render_ms} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/methods/wrapper_vdgs.py b/methods/wrapper_vdgs.py new file mode 100644 index 0000000000000000000000000000000000000000..286cf8b98043a733abd8f4458da2c27b5f00ca9c --- /dev/null +++ b/methods/wrapper_vdgs.py @@ -0,0 +1,259 @@ +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +from argparse import ArgumentParser +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.insert(0, '/root/autodl-tmp/VDGS_offy') +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +@register_method("vdgs") +class VDGSWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset.sh_degree) + # INJECTED_RES_FIX begin + import sys as _sys + _scene, _explicit_res = None, None + for _i, _a in enumerate(_sys.argv[:-1]): + _v = _sys.argv[_i + 1] + if _a == "--scene": _scene = _v + elif _a == "--source_path": _scene = _v.rstrip("/").split("/")[-1] + elif _a == "--resolution": + try: _explicit_res = int(_v) + except: pass + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + if _explicit_res is not None and _explicit_res > 0: + _res = _explicit_res + elif _scene is not None: + _res = 4 if _scene in _OUTDOOR_360 else 2 + else: + _res = None + try: + if _res is not None: + self.dataset.resolution = _res + print("[res-fix] scene=%s explicit=%s -> res=%s (%s)" % (_scene, _explicit_res, _res, __file__)) + except Exception as _e: + print("[res-fix] FAILED:", _e) + # INJECTED_RES_FIX end + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + + def _get_u(self, param, grad, opt): + if grad is None: + return None + state = opt.state.get(param) + if state and "exp_avg_sq" in state: + v = state["exp_avg_sq"] + for group in opt.param_groups: + if id(param) in [id(p) for p in group["params"]]: + lr = group["lr"] + return (lr / (torch.sqrt(v) + 1e-8)) * grad + return grad + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + bg = torch.rand((3), device="cuda") if self.opt.random_background else self.background + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, bg) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gt_image = viewpoint_cam.original_image.cuda() + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + loss_parasitic = self.opt.lambda_dssim * (1.0 - ssim_value) + loss = loss_target + loss_parasitic + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + stats = {} + if self.track_decoupling and step % 100 == 0: + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.vdgs_optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + grad_target_xyz = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.vdgs_optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + grad_parasitic_xyz = self.gaussians._xyz.grad.clone() if self.gaussians._xyz.grad is not None else torch.zeros_like(self.gaussians._xyz) + valid_mask = (torch.norm(grad_target_xyz, dim=1) > 0) & (torch.norm(grad_parasitic_xyz, dim=1) > 0) + if valid_mask.any(): + grad_cos_sim = float(F.cosine_similarity(grad_target_xyz[valid_mask], grad_parasitic_xyz[valid_mask], dim=1).mean()) + parasitic_ratio = float(torch.norm(grad_parasitic_xyz, dim=1).mean() / (torch.norm(grad_target_xyz, dim=1).mean() + 1e-7)) + param_groups_map = { + "spatial": ([self.gaussians._xyz], self.gaussians.optimizer), + "geometry": ([self.gaussians._scaling, self.gaussians._rotation], self.gaussians.optimizer), + "opacity": ([self.gaussians._opacity], self.gaussians.optimizer), + "appearance": ([self.gaussians._features_dc, self.gaussians._features_rest], self.gaussians.optimizer) + } + vdgs_params = list(self.gaussians._vdgs.parameters()) if self.gaussians._vdgs else [] + if vdgs_params: + param_groups_map["appearance"] = (param_groups_map["appearance"][0] + vdgs_params, self.gaussians.vdgs_optimizer) + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.vdgs_optimizer.zero_grad(set_to_none=True) + loss_target.backward(retain_graph=True) + u_target = {} + for group_name, (params, opt) in param_groups_map.items(): + u_list = [] + for p in params: + if p.grad is not None: + u_t = self._get_u(p, p.grad.clone(), opt) + if u_t is not None: + u_list.append(u_t.reshape(-1)) + u_target[group_name] = torch.cat(u_list) if u_list else torch.tensor([], device="cuda") + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.vdgs_optimizer.zero_grad(set_to_none=True) + loss_parasitic.backward(retain_graph=True) + u_parasitic = {} + for group_name, (params, opt) in param_groups_map.items(): + u_list = [] + for p in params: + if p.grad is not None: + u_p = self._get_u(p, p.grad.clone(), opt) + if u_p is not None: + u_list.append(u_p.reshape(-1)) + u_parasitic[group_name] = torch.cat(u_list) if u_list else torch.tensor([], device="cuda") + for group_name in param_groups_map.keys(): + gt = u_target.get(group_name) + gp = u_parasitic.get(group_name) + if gt is not None and gp is not None and gt.numel() > 0 and gp.numel() > 0 and gt.norm() > 0 and gp.norm() > 0: + cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0))) + r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7)) + ti = r * max(0.0, -cos) + else: + ti = 0.0 + stats[f"sti_{group_name}"] = ti + self.gaussians.optimizer.zero_grad(set_to_none=True) + self.gaussians.vdgs_optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + loss.backward() + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 20 if step > self.opt.opacity_reset_interval else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold) + if step < self.opt.iterations: + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + if self.gaussians.vdgs_type: + self.gaussians.vdgs_optimizer.step() + self.gaussians.vdgs_optimizer.zero_grad(set_to_none=True) + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), + "loss_l1": float(loss_target), + "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), + "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)) + } + if self.track_decoupling and step % 100 == 0: + metrics["grad_cos_sim"] = float(grad_cos_sim) + metrics["parasitic_ratio"] = float(parasitic_ratio) + metrics.update(stats) + self.last_n_gaussians = num_gaussians + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + scales_2d = scales[:, :2] if scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + gamma = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + ckpt = os.path.join(model_path, f"chkpnt{iteration}.pth") + if os.path.exists(ckpt): + model_params, _ = torch.load(ckpt) + self.gaussians.restore(model_params, self.opt) + else: + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2) + max_S, _ = torch.max(scales_2d, dim=1) + min_S, _ = torch.min(scales_2d, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales_2d)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = (scales[:, :2].max(dim=1)[0].pow(2)) if scales.shape[1] >= 2 else scales.squeeze().pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/aggregate_and_plot.py b/scripts/aggregate_and_plot.py new file mode 100644 index 0000000000000000000000000000000000000000..b9799eb5415dc847672d852c79c20d7de069a278 --- /dev/null +++ b/scripts/aggregate_and_plot.py @@ -0,0 +1,64 @@ +import os +import json +import pandas as pd +import seaborn as sns +import matplotlib.pyplot as plt + +def main(): + data = [] + out_dir = "outputs" + + if not os.path.exists(out_dir): + return + + for d in os.listdir(out_dir): + json_path = os.path.join(out_dir, d, "metrics.json") + if os.path.exists(json_path): + try: + with open(json_path, 'r') as f: + j = json.load(f) + + flags = j.get("pathology_flags", []) + + row = { + "method": j.get("method", "unknown"), + "scene": j.get("scene", "unknown"), + "PSNR": j.get("photometric", {}).get("PSNR", 0.0), + "LPIPS": j.get("photometric", {}).get("LPIPS", 0.0), + "N_gaussians": j.get("geometric", {}).get("N_gaussians", 0), + "gamma_median": j.get("geometric", {}).get("gamma_median", 0.0), + "alpha_mean": j.get("geometric", {}).get("alpha_mean", 0.0), + "is_collapsed": len(flags) > 0, + "flags": ",".join(flags) + } + data.append(row) + except Exception: + pass + + if not data: + return + + df = pd.DataFrame(data) + df.to_csv("SplatAtlas_Full_Report.csv", index=False) + + plt.figure(figsize=(10, 6)) + sns.scatterplot( + data=df, + x="PSNR", + y="gamma_median", + hue="method", + style="is_collapsed", + s=150, + alpha=0.8 + ) + plt.axhline(18.0, color='red', linestyle='--', label='Pathology Threshold (\u03b3=18)') + plt.yscale('log') + plt.title("Photometric-to-Geometric Compensation (Cat 1-5)") + plt.xlabel("PSNR (Higher is better)") + plt.ylabel("Anisotropy \u03b3 (Lower is better, log scale)") + plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left') + plt.tight_layout() + plt.savefig("compensation_dynamics.png", dpi=300) + +if __name__ == "__main__": + main() diff --git a/scripts/aggregate_surfel_hypothesis_all.py b/scripts/aggregate_surfel_hypothesis_all.py new file mode 100644 index 0000000000000000000000000000000000000000..caa6accad0651cef4e907574688b36d543548934 --- /dev/null +++ b/scripts/aggregate_surfel_hypothesis_all.py @@ -0,0 +1,402 @@ +#!/usr/bin/env python3 +import json +from pathlib import Path + +import pandas as pd + + +PROJECT = Path("/root/autodl-tmp/SplatAtlas") +OUT_ROOT = PROJECT / "outputs" / "surfel_hypothesis_all" +SOURCE_JSON_ROOT = PROJECT / "outputs" / "surfel_hypothesis" +RUN_INDEX = OUT_ROOT / "run_index.csv" + +OUT_ROOT.mkdir(parents=True, exist_ok=True) + + +METHOD_GROUP = { + "vanilla_3dgs": "standard_3dgs", + + "3dgs_dr": "standard_3dgs", + "3dgsmcmc": "standard_3dgs", + "aaags": "standard_3dgs_nondefault_checkpoint", + "absgs": "standard_3dgs", + "absgssgf": "standard_3dgs", + "analyticsplatting": "standard_3dgs", + "atomgs": "standard_3dgs", + "cdcgs": "standard_3dgs", + "coadaptgs": "standard_3dgs", + "conegs": "standard_3dgs", + "drkgs": "standard_3dgs", + "erankgs": "standard_3dgs", + "ges": "standard_3dgs", + "ghap": "standard_3dgs", + "gslpm": "standard_3dgs", + "hogs": "standard_3dgs_nondefault_checkpoint", + "lapisgs": "standard_3dgs", + "lightgaussian": "standard_3dgs", + "minisplatting": "standard_3dgs", + "minisplattingsgf": "standard_3dgs", + "mipsplatting": "standard_3dgs", + "opti3dgs": "standard_3dgs", + "pixelgs": "standard_3dgs", + "pixelgssgf": "standard_3dgs", + "rdogaussian": "standard_3dgs", + "reactgs": "standard_3dgs", + "steepgs": "standard_3dgs", + "tamgs": "standard_3dgs_nondefault_checkpoint", + + "gaussianpro": "geometry_aware_standard_or_usable", + "gof": "geometry_aware_standard_or_usable", + "pgsr": "geometry_aware_standard_or_usable", + "gsdf": "geometry_aware_standard_or_usable", + "gspull": "geometry_aware_standard_or_usable", + + "2dgs": "surface_primitive_2dgs_disk_normal", + "trimgs": "surface_primitive_pca_only", + "gaussian_surfel": "gaussian_surfel_usable", + "scaffoldgs": "hierarchical_or_anchor_usable", + "octree_gs": "hierarchical_or_anchor_usable", + "lod_gs": "hierarchical_or_anchor_usable", + "sugar": "surface_mesh_coupled_usable", +} + + +METHODS = list(METHOD_GROUP.keys()) +SCENES = ["barn", "caterpillar", "truck"] + + +def get(d, key): + return d.get(key, None) + + +def read_failure_tail(log_path, n=25): + p = PROJECT / log_path if not str(log_path).startswith("/") else Path(log_path) + if not p.exists(): + return "" + lines = p.read_text(errors="replace").splitlines() + return "\n".join(lines[-n:]) + + +rows = [] +failures = [] + +run_index_df = None +if RUN_INDEX.exists(): + run_index_df = pd.read_csv(RUN_INDEX) +else: + run_index_df = pd.DataFrame(columns=["method", "scene", "group", "status", "log_path", "json_path"]) + + +for method in METHODS: + for scene in SCENES: + p = SOURCE_JSON_ROOT / f"{method}_{scene}" / "surfel_hypothesis_eval.json" + + idx_rows = run_index_df[ + (run_index_df.get("method") == method) & + (run_index_df.get("scene") == scene) + ] + + log_path = None + run_status = "not_in_run_index" + iteration_arg = "" + + if len(idx_rows) > 0: + rr = idx_rows.iloc[-1] + log_path = rr.get("log_path", None) + run_status = rr.get("status", "unknown") + iteration_arg = rr.get("iteration_arg", "") + + if not p.exists(): + failures.append({ + "method": method, + "scene": scene, + "group": METHOD_GROUP[method], + "status": run_status, + "iteration_arg": iteration_arg, + "json_path": str(p), + "log_path": log_path, + "tail": read_failure_tail(log_path) if log_path else "", + }) + rows.append({ + "method": method, + "scene": scene, + "group": METHOD_GROUP[method], + "status": "missing_json", + "iteration_arg": iteration_arg, + "json_path": str(p), + "log_path": log_path, + }) + continue + + try: + d = json.loads(p.read_text()) + except Exception as e: + failures.append({ + "method": method, + "scene": scene, + "group": METHOD_GROUP[method], + "status": f"json_read_error: {type(e).__name__}: {e}", + "iteration_arg": iteration_arg, + "json_path": str(p), + "log_path": log_path, + "tail": read_failure_tail(log_path) if log_path else "", + }) + rows.append({ + "method": method, + "scene": scene, + "group": METHOD_GROUP[method], + "status": "json_read_error", + "iteration_arg": iteration_arg, + "json_path": str(p), + "log_path": log_path, + }) + continue + + row = { + "method": method, + "scene": scene, + "group": METHOD_GROUP[method], + "status": "ok", + "iteration_arg": iteration_arg, + + "json_path": str(p), + "log_path": log_path, + + "ply_path": get(d, "ply_path"), + "n_gaussians_recon": get(d, "n_gaussians_recon"), + "n_recon_after_crop": get(d, "n_recon_after_crop"), + "n_recon_eval": get(d, "n_recon_eval"), + "n_gt_after_crop": get(d, "n_gt_after_crop"), + "n_surface_near": get(d, "n_surface_near"), + "surface_near_ratio": get(d, "surface_near_ratio"), + "n_normal_eval": get(d, "n_normal_eval"), + + "tau": get(d, "tau"), + "surface_near_threshold": get(d, "surface_near_threshold"), + "distance_median": get(d, "distance_recon_to_gt_median"), + "distance_mean": get(d, "distance_recon_to_gt_mean"), + + "has_standard_covariance_shape": get(d, "has_standard_covariance_shape"), + "has_v3_normal": get(d, "has_v3_normal"), + "has_2dgs_normal": get(d, "has_2dgs_normal"), + "has_explicit_normal": get(d, "has_explicit_normal"), + + # Normal metrics. + "center_pca_normal_median_deg": get(d, "center_pca_normal_error_median_deg"), + "center_pca_normal_mean_deg": get(d, "center_pca_normal_error_mean_deg"), + "center_pca_normal_iqr_deg": get(d, "center_pca_normal_error_iqr_deg"), + + "v3_normal_median_deg": get(d, "v3_normal_error_median_deg"), + "v3_normal_mean_deg": get(d, "v3_normal_error_mean_deg"), + "v3_normal_iqr_deg": get(d, "v3_normal_error_iqr_deg"), + "delta_v3_minus_center_pca_deg": get(d, "median_delta_v3_minus_center_pca_deg"), + "center_pca_better_than_v3": get(d, "center_pca_better_than_v3_by_median"), + + "two_dgs_normal_median_deg": get(d, "two_dgs_normal_error_median_deg"), + "two_dgs_normal_mean_deg": get(d, "two_dgs_normal_error_mean_deg"), + "two_dgs_normal_iqr_deg": get(d, "two_dgs_normal_error_iqr_deg"), + "delta_2dgs_minus_center_pca_deg": get(d, "median_delta_2dgs_minus_center_pca_deg"), + "center_pca_better_than_2dgs": get(d, "center_pca_better_than_2dgs_by_median"), + + "explicit_normal_median_deg": get(d, "explicit_normal_error_median_deg"), + "explicit_normal_mean_deg": get(d, "explicit_normal_error_mean_deg"), + "explicit_normal_iqr_deg": get(d, "explicit_normal_error_iqr_deg"), + "delta_explicit_minus_center_pca_deg": get(d, "median_delta_explicit_minus_center_pca_deg"), + + # Surface confidence calibration. + "opacity_auc_surface": get(d, "opacity_sigmoid_auroc_surface_near"), + "opacity_ap_surface": get(d, "opacity_sigmoid_ap_surface_near"), + "opacity_spearman_minus_distance": get(d, "opacity_sigmoid_spearman_vs_minus_distance"), + "opacity_median": get(d, "opacity_sigmoid_median"), + + "cov_log_anisotropy_median": get(d, "cov_log_anisotropy_sigma_median"), + "cov_anisotropy_ratio_median": get(d, "cov_anisotropy_sigma_ratio_median"), + "cov_anisotropy_auc_surface": get(d, "cov_log_anisotropy_sigma_auroc_surface_near"), + "cov_anisotropy_ap_surface": get(d, "cov_log_anisotropy_sigma_ap_surface_near"), + "cov_anisotropy_spearman_minus_distance": get(d, "cov_log_anisotropy_sigma_spearman_vs_minus_distance"), + "cov_anisotropy_spearman_minus_v3_error": get(d, "cov_log_anisotropy_sigma_spearman_vs_minus_v3_error"), + + "cov_planarity_median": get(d, "cov_planarity_median"), + "cov_planarity_auc_surface": get(d, "cov_planarity_auroc_surface_near"), + "cov_planarity_ap_surface": get(d, "cov_planarity_ap_surface_near"), + "cov_planarity_spearman_minus_distance": get(d, "cov_planarity_spearman_vs_minus_distance"), + "cov_planarity_spearman_minus_v3_error": get(d, "cov_planarity_spearman_vs_minus_v3_error"), + + "cov_linearity_median": get(d, "cov_linearity_median"), + "cov_scattering_median": get(d, "cov_scattering_median"), + + "spatial_planarity_median": get(d, "spatial_pca_planarity_median"), + "spatial_planarity_auc_surface": get(d, "spatial_pca_planarity_auroc_surface_near"), + "spatial_planarity_spearman_minus_distance": get(d, "spatial_pca_planarity_spearman_vs_minus_distance"), + "spatial_planarity_spearman_minus_center_pca_error": get(d, "spatial_pca_planarity_spearman_vs_minus_center_pca_error"), + + "spatial_linearity_median": get(d, "spatial_pca_linearity_median"), + "spatial_scattering_median": get(d, "spatial_pca_scattering_median"), + + "gt_local_pca_planarity_median": get(d, "gt_local_pca_planarity_median"), + "gt_local_pca_linearity_median": get(d, "gt_local_pca_linearity_median"), + "gt_local_pca_scattering_median": get(d, "gt_local_pca_scattering_median"), + + "wall_time_seconds": get(d, "wall_time_seconds"), + } + + rows.append(row) + + +df = pd.DataFrame(rows) +fail_df = pd.DataFrame(failures) + +aggregate_csv = OUT_ROOT / "surfel_hypothesis_all_aggregate.csv" +failure_csv = OUT_ROOT / "surfel_hypothesis_all_failures.csv" +method_mean_csv = OUT_ROOT / "surfel_hypothesis_all_method_mean.csv" +scene_mean_csv = OUT_ROOT / "surfel_hypothesis_all_scene_mean.csv" +report_md = OUT_ROOT / "surfel_hypothesis_all_report.md" + +df.to_csv(aggregate_csv, index=False) +fail_df.to_csv(failure_csv, index=False) + +ok = df[df["status"] == "ok"].copy() + +numeric_cols = [ + "surface_near_ratio", + "distance_median", + "center_pca_normal_median_deg", + "v3_normal_median_deg", + "two_dgs_normal_median_deg", + "delta_v3_minus_center_pca_deg", + "delta_2dgs_minus_center_pca_deg", + "opacity_auc_surface", + "cov_anisotropy_auc_surface", + "cov_planarity_auc_surface", + "spatial_planarity_auc_surface", + "cov_anisotropy_spearman_minus_distance", + "cov_planarity_spearman_minus_distance", + "opacity_spearman_minus_distance", +] + +if len(ok): + method_mean = ok.groupby(["group", "method"], dropna=False)[numeric_cols].mean(numeric_only=True).reset_index() + scene_mean = ok.groupby(["scene"], dropna=False)[numeric_cols].mean(numeric_only=True).reset_index() +else: + method_mean = pd.DataFrame() + scene_mean = pd.DataFrame() + +method_mean.to_csv(method_mean_csv, index=False) +scene_mean.to_csv(scene_mean_csv, index=False) + +# Useful derived summaries. +num_expected = len(METHODS) * len(SCENES) +num_ok = int((df["status"] == "ok").sum()) +num_fail = num_expected - num_ok + +best_v3 = ok.dropna(subset=["v3_normal_median_deg"]).sort_values("v3_normal_median_deg").head(15) +best_2dgs = ok.dropna(subset=["two_dgs_normal_median_deg"]).sort_values("two_dgs_normal_median_deg").head(15) +weak_aniso = ok.dropna(subset=["cov_anisotropy_auc_surface"]).copy() + +with open(report_md, "w") as f: + f.write("# Surfel Hypothesis Full Run Report\n\n") + f.write(f"- Expected cells: {num_expected}\n") + f.write(f"- Successful cells: {num_ok}\n") + f.write(f"- Failed/missing cells: {num_fail}\n") + f.write(f"- Aggregate CSV: `{aggregate_csv}`\n") + f.write(f"- Failure CSV: `{failure_csv}`\n") + f.write(f"- Method mean CSV: `{method_mean_csv}`\n") + f.write(f"- Scene mean CSV: `{scene_mean_csv}`\n\n") + + f.write("## Group counts among successful cells\n\n") + if len(ok): + f.write(ok["group"].value_counts().to_frame("count").to_markdown()) + f.write("\n\n") + + f.write("## Best covariance-derived V3 normals, lower is better\n\n") + if len(best_v3): + cols = [ + "method", "scene", "group", + "surface_near_ratio", + "v3_normal_median_deg", + "v3_normal_iqr_deg", + "center_pca_normal_median_deg", + "delta_v3_minus_center_pca_deg", + ] + f.write(best_v3[cols].to_markdown(index=False)) + else: + f.write("No V3 normal results.\n") + f.write("\n\n") + + f.write("## 2DGS / disk normal results\n\n") + if len(best_2dgs): + cols = [ + "method", "scene", "group", + "surface_near_ratio", + "two_dgs_normal_median_deg", + "two_dgs_normal_iqr_deg", + "center_pca_normal_median_deg", + "delta_2dgs_minus_center_pca_deg", + ] + f.write(best_2dgs[cols].to_markdown(index=False)) + else: + f.write("No 2DGS disk-normal results.\n") + f.write("\n\n") + + f.write("## Mean by method\n\n") + if len(method_mean): + show_cols = [ + "group", "method", + "surface_near_ratio", + "v3_normal_median_deg", + "two_dgs_normal_median_deg", + "center_pca_normal_median_deg", + "cov_anisotropy_auc_surface", + "cov_planarity_auc_surface", + "opacity_auc_surface", + ] + existing = [c for c in show_cols if c in method_mean.columns] + f.write(method_mean[existing].to_markdown(index=False)) + f.write("\n\n") + + f.write("## Failures\n\n") + if len(fail_df): + f.write(fail_df[["method", "scene", "group", "status", "log_path"]].to_markdown(index=False)) + else: + f.write("No failures.\n") + f.write("\n") + +print("============================================================") +print("Aggregation done") +print("============================================================") +print(f"[WROTE] {aggregate_csv}") +print(f"[WROTE] {failure_csv}") +print(f"[WROTE] {method_mean_csv}") +print(f"[WROTE] {scene_mean_csv}") +print(f"[WROTE] {report_md}") + +print("\n===== successful / expected =====") +print(num_ok, "/", num_expected) + +print("\n===== failures =====") +if len(fail_df): + print(fail_df[["method", "scene", "group", "status", "log_path"]].to_string(index=False)) +else: + print("No failures.") + +print("\n===== best V3 normals =====") +if len(best_v3): + print(best_v3[[ + "method", "scene", "surface_near_ratio", + "v3_normal_median_deg", + "center_pca_normal_median_deg", + "delta_v3_minus_center_pca_deg", + ]].to_string(index=False)) +else: + print("No V3 results.") + +print("\n===== 2DGS / disk normal results =====") +if len(best_2dgs): + print(best_2dgs[[ + "method", "scene", "surface_near_ratio", + "two_dgs_normal_median_deg", + "center_pca_normal_median_deg", + "delta_2dgs_minus_center_pca_deg", + ]].to_string(index=False)) +else: + print("No 2DGS disk normal results.") diff --git a/scripts/analyze_densification_events.py b/scripts/analyze_densification_events.py new file mode 100644 index 0000000000000000000000000000000000000000..abf9006ae65c64ae238da3b520d20a94ed1e4eed --- /dev/null +++ b/scripts/analyze_densification_events.py @@ -0,0 +1,49 @@ +"""Aggregate densification_events.json across scenes/methods, emit Table 8 data.""" +import os +import json +import argparse + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--scenes", nargs="+", + default=["francis", "m60", "panther", "bonsai", "bicycle"]) + ap.add_argument("--methods", nargs="+", + default=["vanilla_logged", "vanillasgf_logged"]) + ap.add_argument("--outputs_root", default="/root/autodl-tmp/SplatAtlas/outputs") + args = ap.parse_args() + + print(f"{'Method':<22} {'Scene':<10} {'#Events':>9} {'#Edge':>8} {'#Flat':>8} {'Edge/Flat':>10}") + print("-" * 72) + + summary = {} + for method in args.methods: + agg_total = agg_edge = agg_flat = 0 + for scene in args.scenes: + log_path = os.path.join(args.outputs_root, f"{method}_{scene}", + "densification_events.json") + if not os.path.exists(log_path): + print(f"{method:<22} {scene:<10} {'(missing)':>9}") + continue + with open(log_path) as f: + events = json.load(f) + total = sum(e["n_will_densify"] for e in events) + edge = sum(e["n_edge"] for e in events) + flat = sum(e["n_flat"] for e in events) + ratio = edge / max(flat, 1) + agg_total += total; agg_edge += edge; agg_flat += flat + print(f"{method:<22} {scene:<10} {total:>9} {edge:>8} {flat:>8} {ratio:>10.2f}") + m_ratio = agg_edge / max(agg_flat, 1) + print(f"{method:<22} {'AVERAGE':<10} {agg_total:>9} {agg_edge:>8} {agg_flat:>8} {m_ratio:>10.2f}") + print("-" * 72) + summary[method] = {"total": agg_total, "edge": agg_edge, "flat": agg_flat, "ratio": m_ratio} + + if "vanilla_logged" in summary and "vanillasgf_logged" in summary: + base_r = summary["vanilla_logged"]["ratio"] + sgf_r = summary["vanillasgf_logged"]["ratio"] + print(f"\nEdge/Flat ratio shift: baseline {base_r:.2f}x → SGF {sgf_r:.2f}x" + f" (Δ = {sgf_r - base_r:+.2f})") + + +if __name__ == "__main__": + main() diff --git a/scripts/analyze_gradient_intensity.py b/scripts/analyze_gradient_intensity.py new file mode 100644 index 0000000000000000000000000000000000000000..a7581ec114a338fac2ca3020c8a3060337327748 --- /dev/null +++ b/scripts/analyze_gradient_intensity.py @@ -0,0 +1,44 @@ +"""Aggregate gradient_intensity.json across scenes/methods → Table 5 expansion.""" +import os, json, argparse +import numpy as as np if False else None +import numpy as np + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--methods", nargs="+", default=["vanilla_logged", "vanillasgf_logged"]) + ap.add_argument("--scenes", nargs="+", default=["francis","m60","panther","bonsai","bicycle","garden"]) + ap.add_argument("--outputs_root", default="/root/autodl-tmp/SplatAtlas/outputs") + args = ap.parse_args() + + print(f"{'Method':<22} {'Scene':<10} {'Edge grad':>11} {'Flat grad':>11} {'Edge/Flat':>10}") + print("-"*68) + summary = {} + for method in args.methods: + ae, af = [], [] + for scene in args.scenes: + log_path = os.path.join(args.outputs_root, f"{method}_{scene}", "gradient_intensity.json") + if not os.path.exists(log_path): + print(f"{method:<22} {scene:<10} (missing)"); continue + with open(log_path) as f: entries = json.load(f) + ev = [e["edge_mean"] for e in entries if e.get("edge_mean", 0) > 0] + fv = [e["flat_mean"] for e in entries if e.get("flat_mean", 0) > 0] + if not ev or not fv: + print(f"{method:<22} {scene:<10} (empty)"); continue + em, fm = float(np.mean(ev)), float(np.mean(fv)) + r = em / max(fm, 1e-9) + print(f"{method:<22} {scene:<10} {em:>11.3e} {fm:>11.3e} {r:>10.2f}") + ae.append(em); af.append(fm) + if ae and af: + mep, mfp = float(np.mean(ae)), float(np.mean(af)) + mr = mep / max(mfp, 1e-9) + print(f"{method:<22} {'AVERAGE':<10} {mep:>11.3e} {mfp:>11.3e} {mr:>10.2f}") + print("-"*68) + summary[method] = {"edge": mep, "flat": mfp, "ratio": mr} + if "vanilla_logged" in summary and "vanillasgf_logged" in summary: + print(f"\nEdge/Flat grad ratio shift: baseline {summary['vanilla_logged']['ratio']:.2f}x" + f" → SGF {summary['vanillasgf_logged']['ratio']:.2f}x") + + +if __name__ == "__main__": + main() diff --git a/scripts/audit_submission.py b/scripts/audit_submission.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/build_grand_table.py b/scripts/build_grand_table.py new file mode 100644 index 0000000000000000000000000000000000000000..09ad054adff245b721d0f8d9a7c34c4906d76522 --- /dev/null +++ b/scripts/build_grand_table.py @@ -0,0 +1,65 @@ +import os +import json +import pandas as pd + +def main(): + outputs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "outputs")) + data = [] + + for folder in os.listdir(outputs_dir): + method_dir = os.path.join(outputs_dir, folder) + if not os.path.isdir(method_dir) or folder == "pathology_archive": + continue + + # 优先寻找批量管线的终点文件 + metrics_file = os.path.join(method_dir, "metrics_test_iter30000.json") + if not os.path.exists(metrics_file): + metrics_file = os.path.join(method_dir, "metrics.json") + + train_metrics_file = os.path.join(method_dir, "metrics_train_iter30000.json") + physics_file = os.path.join(method_dir, "offline_physics_30000.json") + + row = {"Method": folder.split('_')[0], "Scene": folder.split('_')[1] if '_' in folder else "unknown"} + + if os.path.exists(metrics_file): + with open(metrics_file, 'r') as f: + m = json.load(f) + row["PSNR_Test"] = m.get("photometric", {}).get("PSNR", 0) + row["SSIM"] = m.get("photometric", {}).get("SSIM", 0) + row["Prior"] = m.get("prior_type", "none") + row["Flags"] = len(m.get("pathology_flags", [])) + + # 提取 Train PSNR 并计算你的核心主张：Delta PSNR + if os.path.exists(train_metrics_file): + with open(train_metrics_file, 'r') as tf: + tm = json.load(tf) + row["PSNR_Train"] = tm.get("photometric", {}).get("PSNR", 0) + row["Delta_PSNR"] = row["PSNR_Train"] - row["PSNR_Test"] + + if os.path.exists(physics_file): + with open(physics_file, 'r') as f: + p = json.load(f) + # 修正了旧版错误的字典查找，完美对齐新数据 + row["Gini"] = p.get("spatial_gini", 0) + row["SH_Ratio"] = p.get("sh_energy_ratio", 0) + row["Condition"] = p.get("scale_condition_mean", 0) + row["Billboard_Bias"] = p.get("billboard_bias_ratio", 0) + + if len(row) > 2: + data.append(row) + + if not data: + return + + df = pd.DataFrame(data).round(3) + csv_path = os.path.join(outputs_dir, "grand_leaderboard.csv") + df.to_csv(csv_path, index=False) + + latex_str = df.to_latex(index=False, float_format="%.3f") + with open(os.path.join(outputs_dir, "grand_leaderboard.tex"), "w") as f: + f.write(latex_str) + + print(f"📊 [GrandTable] 大榜单生成完毕！共收录 {len(data)} 个模型数据。") + +if __name__ == "__main__": + main() diff --git a/scripts/check_tb.py b/scripts/check_tb.py new file mode 100644 index 0000000000000000000000000000000000000000..16b044b33bdc7f928d921a88908af9d3f5b64e8b --- /dev/null +++ b/scripts/check_tb.py @@ -0,0 +1,57 @@ + +import os + +from tensorboard.backend.event_processing.event_accumulator import EventAccumulator + + + +target_dir = "/root/autodl-tmp/SplatAtlas/outputs/2dgs_barn" + +if not os.path.exists(target_dir): + + print("Directory not found.") + + exit() + + + +tb_files = [os.path.join(target_dir, f) for f in os.listdir(target_dir) if "events.out.tfevents" in f] + + + +if not tb_files: + + print("No TensorBoard files found.") + + exit() + + + +ea = EventAccumulator(tb_files[0]) + +ea.Reload() + + + +print("Recorded Scalars:") + +for tag in ea.Tags().get('scalars', []): + + print(f" - {tag}") + + + +print("\nRecorded Histograms:") + +for tag in ea.Tags().get('histograms', []): + + print(f" - {tag}") + + + +if 'peak_vram_GB' in ea.Tags().get('scalars', []): + + vram_events = ea.Scalars('peak_vram_GB') + + print(f"\nLast recorded VRAM: {vram_events[-1].value:.2f} GB at step {vram_events[-1].step}") + diff --git a/scripts/check_tb_values.py b/scripts/check_tb_values.py new file mode 100644 index 0000000000000000000000000000000000000000..5087a8bf8a2832334e19b4ca0a3fbf510e16b956 --- /dev/null +++ b/scripts/check_tb_values.py @@ -0,0 +1,36 @@ +import os +from tensorboard.backend.event_processing.event_accumulator import EventAccumulator + +target_dir = "/root/autodl-tmp/SplatAtlas/outputs/2dgs_barn" +tb_files = [os.path.join(target_dir, f) for f in os.listdir(target_dir) if "events.out.tfevents" in f] + +if not tb_files: + print("No events found.") + exit() + +ea = EventAccumulator(tb_files[0]) +ea.Reload() + +# 检查几个关键标量的值 +interesting_scalars = [ + 'train/loss', + 'train/num_gaussians', + 'train/peak_vram_GB', + 'train/grad_cos_sim' +] + +print(f"{'Metric':<25} | {'Last Value':<15} | {'Step':<10}") +print("-" * 55) + +for tag in interesting_scalars: + if tag in ea.Tags().get('scalars', []): + events = ea.Scalars(tag) + if events: + last_event = events[-1] + print(f"{tag:<25} | {last_event.value:<15.6f} | {last_event.step:<10}") + +# 检查直方图是否有数据点 +print("\nHistogram Status:") +for tag in ea.Tags().get('histograms', []): + hist_events = ea.Histograms(tag) + print(f" - {tag:<30}: {len(hist_events)} snapshots recorded") diff --git a/scripts/compute_offline_physics.py b/scripts/compute_offline_physics.py new file mode 100644 index 0000000000000000000000000000000000000000..65ced28dc46ddfe841e613c94157d419bb5bec0b --- /dev/null +++ b/scripts/compute_offline_physics.py @@ -0,0 +1,129 @@ +import os +import json + +class NpEncoder(json.JSONEncoder): + def default(self, obj): + if hasattr(obj, 'item'): + return obj.item() + return super().default(obj) + +import argparse +import torch +import numpy as np +from scipy.stats import entropy +import sys + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from core.registry import METHOD_REGISTRY +import methods + +def compute_gini(array): + array = array.flatten() + array = array[array > 0] + if len(array) == 0: return 0.0 + array = np.sort(array) + index = np.arange(1, array.shape[0] + 1) + n = array.shape[0] + return ((np.sum((2 * index - n - 1) * array)) / (n * np.sum(array))) + +def generate_standard_grid(cameras, resolution=64): + """基于相机分布生成标准查询网格 (默认 64^3 约 26 万点，兼顾精度与显存)""" + if not cameras: + min_bounds = torch.tensor([-5.0, -5.0, -5.0]) + max_bounds = torch.tensor([5.0, 5.0, 5.0]) + else: + cam_centers = torch.stack([c.camera_center for c in cameras]) + min_bounds, _ = torch.min(cam_centers, dim=0) + max_bounds, _ = torch.max(cam_centers, dim=0) + extent = max_bounds - min_bounds + min_bounds -= extent * 0.1 + max_bounds += extent * 0.1 + + X = torch.linspace(min_bounds[0], max_bounds[0], resolution) + Y = torch.linspace(min_bounds[1], max_bounds[1], resolution) + Z = torch.linspace(min_bounds[2], max_bounds[2], resolution) + grid_x, grid_y, grid_z = torch.meshgrid(X, Y, Z, indexing='ij') + + query_points = torch.stack([grid_x.flatten(), grid_y.flatten(), grid_z.flatten()], dim=-1) + return query_points.cuda() + +def analyze_physics(model, cameras): + stats = {} + + # 1. 变体自述物理指标 (特异性指标，如畸变度、Billboard偏差等) + try: + specific_metrics = model.compute_physical_metrics(cameras=cameras) + stats.update(specific_metrics) + except NotImplementedError: + print(" ⚠️ [Physics] Wrapper 未实现 compute_physical_metrics，跳过特异性指标。") + except Exception as e: + print(f" ⚠️ [Physics] 特异性指标计算发生错误: {e}") + + # 2. 宏观标量场评估 (Apple-to-Apple 的绝对公平空间测试) + try: + query_points = generate_standard_grid(cameras, resolution=64) + with torch.no_grad(): + scalar_field = model.evaluate_spatial_field(query_points, cameras=cameras) + + field_np = scalar_field.cpu().numpy() + stats['spatial_gini'] = compute_gini(field_np) + + p = field_np / (field_np.sum() + 1e-7) + stats['density_entropy'] = entropy(p[p > 0]) + stats["Field_Total_Volume"] = float(np.sum(field_np)) + except NotImplementedError: + print(" ⚠️ [Physics] Wrapper 未实现 evaluate_spatial_field，跳过空间分布诊断。") + except Exception as e: + print(f" ⚠️ [Physics] 空间标量场评估发生异常: {e}") + + return {k: round(v, 4) if isinstance(v, float) else v for k, v in stats.items()} + +def load_cameras_from_json(cam_path): + """为了兼容旧流水线，简单构建具有 camera_center 属性的 Dummy Cam""" + if not os.path.exists(cam_path): return [] + with open(cam_path, 'r') as f: cams_data = json.load(f) + + class DummyCam: + def __init__(self, center): + self.camera_center = torch.tensor(center, dtype=torch.float32) + + return [DummyCam(c['center']) for c in cams_data if 'center' in c] + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--method", type=str, required=True, help="模型方法名称") + parser.add_argument("--source_path", type=str, default="", help="原始数据集路径") + parser.add_argument("--model_path", type=str, required=True, help="模型输出目录") + parser.add_argument("--iteration", type=int, default=30000) + args = parser.parse_args() + + model_dir = os.path.abspath(args.model_path) + if not os.path.exists(model_dir): + print(f"⚠️ [Physics] 跳过: 模型目录不存在 {model_dir}") + return + + folder_name = os.path.basename(model_dir) + print(f"🔬 [Physics] 正在精准解剖场景: {folder_name} ...") + + # 实例化对应的 Wrapper 并加载资产 + dataset_config = {"source_path": args.source_path, "model_path": args.model_path, "resolution": 1} + try: + model_class = methods.load_method(args.method) + model = model_class(dataset_config, hyperparams={}) + model.load(args.model_path, args.iteration) + except Exception as e: + print(f"❌ [Physics] 实例化或加载模型失败: {e}") + return + + cameras = model.scene.getTrainCameras() if hasattr(model, 'scene') else [] + + stats = analyze_physics(model, cameras) + + if stats: + stats['experiment'] = folder_name + with open(os.path.join(model_dir, f"offline_physics_{args.iteration}.json"), 'w') as f: + json.dump(stats, f, indent=4, cls=NpEncoder) + print(f" ✨ [Physics] {folder_name} 分析完成！提取指标数: {len(stats)}") + +if __name__ == "__main__": + main() diff --git a/scripts/convert_nsvf_to_blender.py b/scripts/convert_nsvf_to_blender.py new file mode 100644 index 0000000000000000000000000000000000000000..ce7c19c8b14c518b2df9b23ab2b0b8f28d8be2ae --- /dev/null +++ b/scripts/convert_nsvf_to_blender.py @@ -0,0 +1,58 @@ +import os +import json +import numpy as np +from PIL import Image + +def convert_scene(scene_path): + print(f"\n⚙️ 转换场景: {scene_path}") + rgb_dir = os.path.join(scene_path, "rgb") + if not os.path.exists(rgb_dir): return + + all_files = sorted([f for f in os.listdir(rgb_dir) if f.endswith(".png")]) + if not all_files: return + + sample_img_path = os.path.join(rgb_dir, all_files[0]) + with Image.open(sample_img_path) as img: + w, h = img.size + + with open(os.path.join(scene_path, "intrinsics.txt"), "r") as f: + line = f.readline().strip() + focal = float(line.split()[0]) + + camera_angle_x = 2 * np.arctan(w / (2 * focal)) + pose_dir = os.path.join(scene_path, "pose") + frames = [] + + for f in all_files: + name = os.path.splitext(f)[0] + pose_path = os.path.join(pose_dir, f"{name}.txt") + if not os.path.exists(pose_path): continue + + c2w = np.loadtxt(pose_path).reshape(4, 4) + + # 🔥 核心修复：OpenCV 到 OpenGL 的坐标系转换！ + # 将 Y 轴和 Z 轴反转，让相机“转过头来”看向物体 + c2w[:, 1:3] *= -1 + + frames.append({ + "file_path": f"./rgb/{name}", + "transform_matrix": c2w.tolist() + }) + + num_train = int(len(frames) * 0.8) + base_json = {"camera_angle_x": camera_angle_x} + + with open(os.path.join(scene_path, "transforms_train.json"), "w") as f: + json.dump({**base_json, "frames": frames[:num_train]}, f, indent=4) + + with open(os.path.join(scene_path, "transforms_test.json"), "w") as f: + json.dump({**base_json, "frames": frames[num_train:]}, f, indent=4) + + print(f" ✅ 成功生成修复了坐标系的 transforms_*.json") + +target_scenes = ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] +base_path = "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF" +for s in target_scenes: + p = os.path.join(base_path, s) + if os.path.exists(p): + convert_scene(p) diff --git a/scripts/eval_surfel_hypothesis.py b/scripts/eval_surfel_hypothesis.py new file mode 100644 index 0000000000000000000000000000000000000000..2232304e63e65d9fc18ac97d972a94c5f0477ca6 --- /dev/null +++ b/scripts/eval_surfel_hypothesis.py @@ -0,0 +1,637 @@ +#!/usr/bin/env python3 +import argparse +import json +import math +import sys +import time +from pathlib import Path + +import numpy as np +from scipy.spatial import cKDTree +from scipy.stats import spearmanr, rankdata + +sys.path.append(str(Path(__file__).resolve().parent)) +import eval_tnt_wrapper as W + + +def sigmoid(x): + x = np.asarray(x, dtype=np.float64) + return 1.0 / (1.0 + np.exp(-x)) + + +def query_tree(tree, points, k=1, batch_size=50000): + ds, inds = [], [] + for s in range(0, len(points), batch_size): + e = min(s + batch_size, len(points)) + try: + d, i = tree.query(points[s:e], k=k, workers=-1) + except TypeError: + d, i = tree.query(points[s:e], k=k) + ds.append(d) + inds.append(i) + return np.concatenate(ds, axis=0), np.concatenate(inds, axis=0) + + +def safe_spearman(x, y): + x = np.asarray(x) + y = np.asarray(y) + m = np.isfinite(x) & np.isfinite(y) + if m.sum() < 10: + return float("nan") + if np.std(x[m]) < 1e-12 or np.std(y[m]) < 1e-12: + return float("nan") + return float(spearmanr(x[m], y[m]).correlation) + + +def auroc_score(labels, scores): + labels = np.asarray(labels).astype(bool) + scores = np.asarray(scores, dtype=np.float64) + m = np.isfinite(scores) + labels = labels[m] + scores = scores[m] + n_pos = int(labels.sum()) + n_neg = int((~labels).sum()) + if n_pos == 0 or n_neg == 0: + return float("nan") + ranks = rankdata(scores) + rank_sum_pos = ranks[labels].sum() + auc = (rank_sum_pos - n_pos * (n_pos + 1) / 2.0) / (n_pos * n_neg) + return float(auc) + + +def average_precision(labels, scores): + labels = np.asarray(labels).astype(bool) + scores = np.asarray(scores, dtype=np.float64) + m = np.isfinite(scores) + labels = labels[m] + scores = scores[m] + n_pos = int(labels.sum()) + if n_pos == 0: + return float("nan") + order = np.argsort(-scores) + y = labels[order] + tp = np.cumsum(y) + precision = tp / (np.arange(len(y)) + 1) + ap = precision[y].mean() + return float(ap) + + +def quantile_bin_stats(feature, surface_mask, distance, normal_error=None, n_bins=5): + feature = np.asarray(feature, dtype=np.float64) + surface_mask = np.asarray(surface_mask).astype(bool) + distance = np.asarray(distance, dtype=np.float64) + + valid = np.isfinite(feature) & np.isfinite(distance) + if valid.sum() < n_bins * 10: + return {} + + x = feature[valid] + sm = surface_mask[valid] + dist = distance[valid] + + qs = np.linspace(0, 1, n_bins + 1) + edges = np.quantile(x, qs) + edges[0] = -np.inf + edges[-1] = np.inf + + out = {} + for b in range(n_bins): + lo, hi = edges[b], edges[b + 1] + if b == n_bins - 1: + idx = (x >= lo) & (x <= hi) + else: + idx = (x >= lo) & (x < hi) + + key = f"bin{b+1}" + if idx.sum() == 0: + continue + + out[f"{key}_n"] = int(idx.sum()) + out[f"{key}_feature_median"] = float(np.median(x[idx])) + out[f"{key}_surface_near_ratio"] = float(sm[idx].mean()) + out[f"{key}_distance_median"] = float(np.median(dist[idx])) + + if normal_error is not None: + ne = np.asarray(normal_error, dtype=np.float64) + if len(ne) == len(feature): + ne = ne[valid] + for b in range(n_bins): + lo, hi = edges[b], edges[b + 1] + idx = (x >= lo) & (x <= hi) if b == n_bins - 1 else (x >= lo) & (x < hi) + key = f"bin{b+1}" + if idx.sum() > 0 and np.isfinite(ne[idx]).sum() > 0: + out[f"{key}_normal_error_median"] = float(np.nanmedian(ne[idx])) + + return out + + +def pca_normals_and_shape(points, neighbor_indices, batch_size=50000): + """ + Returns: + normals: eigenvector of smallest eigenvalue + linearity, planarity, scattering from local point distribution + """ + if neighbor_indices.ndim != 2: + raise ValueError("neighbor_indices must be [M, K]") + + M, K = neighbor_indices.shape + normals = np.empty((M, 3), dtype=np.float32) + linearity = np.empty(M, dtype=np.float32) + planarity = np.empty(M, dtype=np.float32) + scattering = np.empty(M, dtype=np.float32) + + for s in range(0, M, batch_size): + e = min(s + batch_size, M) + neigh = points[neighbor_indices[s:e]] + centered = neigh - neigh.mean(axis=1, keepdims=True) + cov = np.einsum("bki,bkj->bij", centered, centered) / max(K - 1, 1) + + vals, vecs = np.linalg.eigh(cov) + # eigh ascending: vals[:,0] <= vals[:,1] <= vals[:,2] + l3 = np.maximum(vals[:, 0], 0.0) + l2 = np.maximum(vals[:, 1], 0.0) + l1 = np.maximum(vals[:, 2], 1e-30) + + normals[s:e] = vecs[:, :, 0].astype(np.float32) + linearity[s:e] = ((l1 - l2) / l1).astype(np.float32) + planarity[s:e] = ((l2 - l3) / l1).astype(np.float32) + scattering[s:e] = (l3 / l1).astype(np.float32) + + return normals, linearity, planarity, scattering + + +def covariance_shape_features(v, names, indices): + required = ["scale_0", "scale_1", "scale_2"] + if not all(k in names for k in required): + return None + + scales_log = np.stack( + [v["scale_0"][indices], v["scale_1"][indices], v["scale_2"][indices]], + axis=1, + ).astype(np.float64) + + # PLY stores log std in standard 3DGS. + sigma = np.exp(scales_log) + lamb = sigma ** 2 + lamb_sort = np.sort(lamb, axis=1)[:, ::-1] # λ1 >= λ2 >= λ3 + + l1 = np.maximum(lamb_sort[:, 0], 1e-30) + l2 = np.maximum(lamb_sort[:, 1], 0.0) + l3 = np.maximum(lamb_sort[:, 2], 0.0) + + line = (l1 - l2) / l1 + plane = (l2 - l3) / l1 + scatter = l3 / l1 + + log_aniso_sigma = np.max(scales_log, axis=1) - np.min(scales_log, axis=1) + aniso_sigma_ratio = np.exp(log_aniso_sigma) + + return { + "cov_linearity": line.astype(np.float32), + "cov_planarity": plane.astype(np.float32), + "cov_scattering": scatter.astype(np.float32), + "cov_log_anisotropy_sigma": log_aniso_sigma.astype(np.float32), + "cov_anisotropy_sigma_ratio": aniso_sigma_ratio.astype(np.float32), + } + + +def gaussian_v3_normals(v, names, indices, transform_mat): + required = ["scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"] + if not all(k in names for k in required): + return None + + idx = indices + scales = np.stack( + [v["scale_0"][idx], v["scale_1"][idx], v["scale_2"][idx]], + axis=1, + ).astype(np.float64) + min_axis = np.argmin(scales, axis=1) + + q = np.stack( + [v["rot_0"][idx], v["rot_1"][idx], v["rot_2"][idx], v["rot_3"][idx]], + axis=1, + ).astype(np.float64) + q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-12) + + # 3DGS convention: [w, x, y, z] + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + + r00 = 1 - 2 * (y * y + z * z) + r01 = 2 * (x * y - w * z) + r02 = 2 * (x * z + w * y) + + r10 = 2 * (x * y + w * z) + r11 = 1 - 2 * (x * x + z * z) + r12 = 2 * (y * z - w * x) + + r20 = 2 * (x * z - w * y) + r21 = 2 * (y * z + w * x) + r22 = 1 - 2 * (x * x + y * y) + + normals = np.empty((len(idx), 3), dtype=np.float64) + + m0 = min_axis == 0 + normals[m0, 0] = r00[m0] + normals[m0, 1] = r10[m0] + normals[m0, 2] = r20[m0] + + m1 = min_axis == 1 + normals[m1, 0] = r01[m1] + normals[m1, 1] = r11[m1] + normals[m1, 2] = r21[m1] + + m2 = min_axis == 2 + normals[m2, 0] = r02[m2] + normals[m2, 1] = r12[m2] + normals[m2, 2] = r22[m2] + + A = transform_mat[:3, :3].astype(np.float64) + normals = normals @ A.T + normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12) + return normals.astype(np.float32) + + + +def two_dgs_normals(v, names, indices, transform_mat): + """ + 2DGS-specific normal. + + 2D Gaussian has two in-plane scale axes, scale_0 and scale_1. + The surface normal is the third local axis of the quaternion rotation, + i.e. the axis perpendicular to the two scaled tangent directions. + + This is NOT the generic 3DGS V3/minor-scale normal. + """ + required = ["scale_0", "scale_1", "rot_0", "rot_1", "rot_2", "rot_3"] + if not all(k in names for k in required): + return None + + # If scale_2 exists, this is probably a standard 3D Gaussian, not 2DGS. + # Let generic V3 handle it instead. + if "scale_2" in names: + return None + + idx = indices + q = np.stack( + [v["rot_0"][idx], v["rot_1"][idx], v["rot_2"][idx], v["rot_3"][idx]], + axis=1, + ).astype(np.float64) + + q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-12) + + # 3DGS/2DGS convention is normally quaternion [w, x, y, z]. + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + + r02 = 2 * (x * z + w * y) + r12 = 2 * (y * z - w * x) + r22 = 1 - 2 * (x * x + y * y) + + # Third column of R: local z axis, perpendicular to the two in-plane axes. + normals = np.stack([r02, r12, r22], axis=1) + + A = transform_mat[:3, :3].astype(np.float64) + normals = normals @ A.T + normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12) + + return normals.astype(np.float32) + + +def explicit_normals(v, names, indices, transform_mat): + candidates = [ + ("nx", "ny", "nz"), + ("normal_x", "normal_y", "normal_z"), + ] + found = None + for a, b, c in candidates: + if a in names and b in names and c in names: + found = (a, b, c) + break + if found is None: + return None + + a, b, c = found + n = np.stack([v[a][indices], v[b][indices], v[c][indices]], axis=1).astype(np.float64) + + raw_norm = np.linalg.norm(n, axis=1) + valid_ratio = float(np.mean(raw_norm > 1e-8)) + + # Many 3DGS-format PLYs contain nx/ny/nz as all-zero placeholders. + # Treat them as absent, otherwise they produce a fake constant 90-degree error. + if valid_ratio < 0.5 or np.nanmedian(raw_norm) < 1e-8: + return None + + A = transform_mat[:3, :3].astype(np.float64) + n = n @ A.T + n = n / (np.linalg.norm(n, axis=1, keepdims=True) + 1e-12) + return n.astype(np.float32) + + +def angular_error(pred, ref): + dots = np.sum(pred * ref, axis=1) + dots = np.clip(np.abs(dots), 0.0, 1.0) + return np.degrees(np.arccos(dots)).astype(np.float32) + + +def angle_stats(prefix, angles): + return { + f"{prefix}_mean_deg": float(np.mean(angles)), + f"{prefix}_median_deg": float(np.median(angles)), + f"{prefix}_q25_deg": float(np.quantile(angles, 0.25)), + f"{prefix}_q75_deg": float(np.quantile(angles, 0.75)), + f"{prefix}_iqr_deg": float(np.quantile(angles, 0.75) - np.quantile(angles, 0.25)), + } + + +def feature_summary(prefix, x): + x = np.asarray(x, dtype=np.float64) + x = x[np.isfinite(x)] + if len(x) == 0: + return {} + return { + f"{prefix}_mean": float(np.mean(x)), + f"{prefix}_median": float(np.median(x)), + f"{prefix}_q25": float(np.quantile(x, 0.25)), + f"{prefix}_q75": float(np.quantile(x, 0.75)), + } + + +def add_feature_diagnostics(result, name, feature, surface_mask, distances): + feature = np.asarray(feature, dtype=np.float64) + result.update(feature_summary(name, feature)) + result[f"{name}_spearman_vs_minus_distance"] = safe_spearman(feature, -distances) + result[f"{name}_auroc_surface_near"] = auroc_score(surface_mask, feature) + result[f"{name}_ap_surface_near"] = average_precision(surface_mask, feature) + + bins = quantile_bin_stats(feature, surface_mask, distances, n_bins=5) + for k, v in bins.items(): + result[f"{name}_{k}"] = v + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--method", required=True) + ap.add_argument("--scene", required=True, choices=sorted(W.SCENE_MAP.keys())) + ap.add_argument("--project-root", default="/root/autodl-tmp/SplatAtlas") + ap.add_argument("--outputs-root", default=None) + ap.add_argument("--tnt-eval-root", default=None) + ap.add_argument("--iteration", type=int, default=None) + + ap.add_argument("--mode", choices=["all", "subsample"], default="subsample") + ap.add_argument("--n-sample", type=int, default=200000) + ap.add_argument("--seed", type=int, default=0) + + ap.add_argument("--distance-multiplier", type=float, default=2.0) + ap.add_argument("--gt-normal-k", type=int, default=30) + ap.add_argument("--recon-pca-k", type=int, default=30) + ap.add_argument("--max-normal-points", type=int, default=50000) + ap.add_argument("--batch-size", type=int, default=50000) + ap.add_argument("--verbose", action="store_true") + args = ap.parse_args() + + t0 = time.time() + + method_in = args.method + if method_in in ["3dgs", "vanilla", "vanilla"]: + method = "vanilla_3dgs" + else: + method = method_in + + project_root = Path(args.project_root) + outputs_root = Path(args.outputs_root) if args.outputs_root else project_root / "outputs" + tnt_eval_root = Path(args.tnt_eval_root) if args.tnt_eval_root else project_root / "data" / "tnt_eval" + + scene = args.scene.lower() + official_scene = W.SCENE_MAP[scene] + tau = W.TAU_DICT[scene] + near_threshold = args.distance_multiplier * tau + + ply_path = W.locate_recon_ply(outputs_root, method, scene, args.iteration) + + scene_eval_dir = tnt_eval_root / official_scene + gt_ply_path = scene_eval_dir / f"{official_scene}.ply" + crop_path = scene_eval_dir / f"{official_scene}.json" + trans_path = scene_eval_dir / f"{official_scene}_trans.txt" + + trans = W.read_transform(trans_path) + crop = W.load_crop(crop_path) + + if args.verbose: + print("=" * 80) + print("Surfel hypothesis test") + print("method:", method) + print("scene:", scene, "->", official_scene) + print("tau:", tau, "surface-near:", near_threshold) + print("recon:", ply_path) + + recon_raw, recon_vertex, recon_names = W.load_vertex_data(ply_path) + recon_aligned = W.apply_transform(recon_raw, trans) + recon_crop_mask = W.crop_mask_tnt(recon_aligned, crop) + recon_crop = recon_aligned[recon_crop_mask] + recon_crop_raw_idx = np.where(recon_crop_mask)[0].astype(np.int64) + + if len(recon_crop) == 0: + raise RuntimeError("Reconstruction crop is empty.") + + eval_idx_in_crop = W.choose_eval_indices( + len(recon_crop), args.mode, args.n_sample, args.seed + ) + recon_eval = recon_crop[eval_idx_in_crop] + recon_eval_raw_idx = recon_crop_raw_idx[eval_idx_in_crop] + + gt_raw, _, _ = W.load_vertex_data(gt_ply_path) + gt_crop_mask = W.crop_mask_tnt(gt_raw, crop) + gt_crop = gt_raw[gt_crop_mask] + if len(gt_crop) == 0: + raise RuntimeError("GT crop is empty.") + + gt_tree = cKDTree(gt_crop) + d_r2g, nn_gt_idx = query_tree(gt_tree, recon_eval, k=1, batch_size=args.batch_size) + + surface_mask = d_r2g < near_threshold + surface_eval_idx = np.where(surface_mask)[0].astype(np.int64) + + rng = np.random.default_rng(args.seed) + if len(surface_eval_idx) > args.max_normal_points: + normal_eval_idx = np.sort( + rng.choice(surface_eval_idx, size=args.max_normal_points, replace=False) + ).astype(np.int64) + else: + normal_eval_idx = surface_eval_idx + + if args.verbose: + print("n_gaussians_recon:", len(recon_raw)) + print("n_recon_after_crop:", len(recon_crop)) + print("n_recon_eval:", len(recon_eval)) + print("n_gt_after_crop:", len(gt_crop)) + print("n_surface_near:", len(surface_eval_idx), "/", len(recon_eval)) + print("n_normal_eval:", len(normal_eval_idx)) + + result = { + "method": method, + "method_input": method_in, + "scene": scene, + "official_scene": official_scene, + "eval_protocol": "surfel_hypothesis_test_tnt_dense_gt_pca_proxy", + "is_official_tnt_metric": False, + "ply_path": str(ply_path), + "gt_ply_path": str(gt_ply_path), + "crop_path": str(crop_path), + "trans_path": str(trans_path), + "tau": float(tau), + "surface_near_threshold": float(near_threshold), + "surface_near_rule": f"d_recon_to_gt < {args.distance_multiplier} * tau", + "mode": args.mode, + "n_sample_requested": int(args.n_sample), + "seed": int(args.seed), + "gt_normal_k": int(args.gt_normal_k), + "recon_pca_k": int(args.recon_pca_k), + "max_normal_points": int(args.max_normal_points), + "n_gaussians_recon": int(len(recon_raw)), + "n_recon_after_crop": int(len(recon_crop)), + "n_recon_eval": int(len(recon_eval)), + "n_gt_after_crop": int(len(gt_crop)), + "n_surface_near": int(len(surface_eval_idx)), + "surface_near_ratio": float(len(surface_eval_idx) / max(len(recon_eval), 1)), + "n_normal_eval": int(len(normal_eval_idx)), + "distance_recon_to_gt_mean": float(np.mean(d_r2g)), + "distance_recon_to_gt_median": float(np.median(d_r2g)), + "distance_recon_to_gt_q25": float(np.quantile(d_r2g, 0.25)), + "distance_recon_to_gt_q75": float(np.quantile(d_r2g, 0.75)), + } + + # Feature 1: opacity. + if "opacity" in recon_names: + opacity = sigmoid(recon_vertex["opacity"][recon_eval_raw_idx]) + add_feature_diagnostics(result, "opacity_sigmoid", opacity, surface_mask, d_r2g) + + # Feature 2: covariance shape, only for standard 3D Gaussian PLYs. + cov_features = covariance_shape_features(recon_vertex, recon_names, recon_eval_raw_idx) + if cov_features is not None: + result["has_standard_covariance_shape"] = True + for name, feat in cov_features.items(): + add_feature_diagnostics(result, name, feat, surface_mask, d_r2g) + else: + result["has_standard_covariance_shape"] = False + + # Feature 3: local spatial PCA shape over Gaussian centers, for all methods. + recon_tree = cKDTree(recon_crop) + _, recon_neighbor_idx_all = query_tree( + recon_tree, recon_eval, k=args.recon_pca_k, batch_size=args.batch_size + ) + recon_pca_normals_all, spatial_linearity, spatial_planarity, spatial_scattering = pca_normals_and_shape( + recon_crop, recon_neighbor_idx_all, batch_size=args.batch_size + ) + + add_feature_diagnostics(result, "spatial_pca_linearity", spatial_linearity, surface_mask, d_r2g) + add_feature_diagnostics(result, "spatial_pca_planarity", spatial_planarity, surface_mask, d_r2g) + add_feature_diagnostics(result, "spatial_pca_scattering", spatial_scattering, surface_mask, d_r2g) + + # Normal consistency on surface-near subset. + if len(normal_eval_idx) > 0: + chosen_points = recon_eval[normal_eval_idx] + chosen_raw_idx = recon_eval_raw_idx[normal_eval_idx] + nearest_gt_points = gt_crop[nn_gt_idx[normal_eval_idx]] + + _, gt_neighbor_idx = query_tree( + gt_tree, nearest_gt_points, k=args.gt_normal_k, batch_size=args.batch_size + ) + gt_normals, gt_linearity, gt_planarity, gt_scattering = pca_normals_and_shape( + gt_crop, gt_neighbor_idx, batch_size=args.batch_size + ) + + # Center PCA normals already computed for all eval points. + center_pca_normals = recon_pca_normals_all[normal_eval_idx] + center_pca_err = angular_error(center_pca_normals, gt_normals) + result.update(angle_stats("center_pca_normal_error", center_pca_err)) + + # V3 normals if standard. + v3 = gaussian_v3_normals(recon_vertex, recon_names, chosen_raw_idx, trans) + if v3 is not None: + v3_err = angular_error(v3, gt_normals) + result["has_v3_normal"] = True + result.update(angle_stats("v3_normal_error", v3_err)) + result["median_delta_v3_minus_center_pca_deg"] = float( + np.median(v3_err) - np.median(center_pca_err) + ) + result["center_pca_better_than_v3_by_median"] = bool( + np.median(center_pca_err) < np.median(v3_err) + ) + else: + result["has_v3_normal"] = False + + # 2DGS-specific normal: perpendicular to the two in-plane scale axes. + n2d = two_dgs_normals(recon_vertex, recon_names, chosen_raw_idx, trans) + if n2d is not None: + n2d_err = angular_error(n2d, gt_normals) + result["has_2dgs_normal"] = True + result.update(angle_stats("two_dgs_normal_error", n2d_err)) + result["median_delta_2dgs_minus_center_pca_deg"] = float( + np.median(n2d_err) - np.median(center_pca_err) + ) + result["center_pca_better_than_2dgs_by_median"] = bool( + np.median(center_pca_err) < np.median(n2d_err) + ) + else: + result["has_2dgs_normal"] = False + + # Explicit normal if available, e.g. 2DGS / surfel / PGSR / GOF exports. + expn = explicit_normals(recon_vertex, recon_names, chosen_raw_idx, trans) + if expn is not None: + exp_err = angular_error(expn, gt_normals) + result["has_explicit_normal"] = True + result.update(angle_stats("explicit_normal_error", exp_err)) + result["median_delta_explicit_minus_center_pca_deg"] = float( + np.median(exp_err) - np.median(center_pca_err) + ) + result["center_pca_better_than_explicit_by_median"] = bool( + np.median(center_pca_err) < np.median(exp_err) + ) + else: + result["has_explicit_normal"] = False + + # Does anisotropy / planarity predict normal reliability? + norm_dist = d_r2g[normal_eval_idx] + dummy_surface = np.ones_like(norm_dist, dtype=bool) + + if cov_features is not None: + for name, feat_all in cov_features.items(): + feat = feat_all[normal_eval_idx] + if result.get("has_v3_normal", False): + result[f"{name}_spearman_vs_minus_v3_error"] = safe_spearman(feat, -v3_err) + result[f"{name}_spearman_vs_minus_center_pca_error"] = safe_spearman(feat, -center_pca_err) + + spatial_feats = { + "spatial_pca_linearity": spatial_linearity[normal_eval_idx], + "spatial_pca_planarity": spatial_planarity[normal_eval_idx], + "spatial_pca_scattering": spatial_scattering[normal_eval_idx], + } + for name, feat in spatial_feats.items(): + if result.get("has_v3_normal", False): + result[f"{name}_spearman_vs_minus_v3_error"] = safe_spearman(feat, -v3_err) + if result.get("has_explicit_normal", False): + result[f"{name}_spearman_vs_minus_explicit_error"] = safe_spearman(feat, -exp_err) + result[f"{name}_spearman_vs_minus_center_pca_error"] = safe_spearman(feat, -center_pca_err) + + result.update(feature_summary("gt_local_pca_planarity", gt_planarity)) + result.update(feature_summary("gt_local_pca_linearity", gt_linearity)) + result.update(feature_summary("gt_local_pca_scattering", gt_scattering)) + + else: + result["has_v3_normal"] = False + result["has_explicit_normal"] = False + result["normal_eval_skipped_reason"] = "No surface-near Gaussians." + + result["wall_time_seconds"] = float(time.time() - t0) + + out_dir = Path(outputs_root) / "surfel_hypothesis" / f"{method}_{scene}" + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "surfel_hypothesis_eval.json" + with open(out_path, "w") as f: + json.dump(result, f, indent=2, sort_keys=True) + + print(json.dumps(result, indent=2, sort_keys=True)) + print(f"\n[WROTE] {out_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_tnt_normals.py b/scripts/eval_tnt_normals.py new file mode 100644 index 0000000000000000000000000000000000000000..a94410f8577e81a24467af385dc15e3130d39e92 --- /dev/null +++ b/scripts/eval_tnt_normals.py @@ -0,0 +1,335 @@ +#!/usr/bin/env python3 +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np +from scipy.spatial import cKDTree + +sys.path.append(str(Path(__file__).resolve().parent)) +import eval_tnt_wrapper as W + + +def sigmoid(x): + return 1.0 / (1.0 + np.exp(-x)) + + +def query_tree(tree, points, k=1, batch_size=200000): + all_d = [] + all_i = [] + for s in range(0, len(points), batch_size): + e = min(s + batch_size, len(points)) + try: + d, i = tree.query(points[s:e], k=k, workers=-1) + except TypeError: + d, i = tree.query(points[s:e], k=k) + all_d.append(d) + all_i.append(i) + return np.concatenate(all_d, axis=0), np.concatenate(all_i, axis=0) + + +def gaussian_normals_from_vertex(v, names, indices, transform_mat): + required = ["scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"] + missing = [k for k in required if k not in names] + if missing: + raise ValueError(f"reconstruction PLY missing Gaussian fields: {missing}") + + idx = indices + + scales = np.stack( + [v["scale_0"][idx], v["scale_1"][idx], v["scale_2"][idx]], + axis=1, + ).astype(np.float64) + + min_axis = np.argmin(scales, axis=1) + + q = np.stack( + [v["rot_0"][idx], v["rot_1"][idx], v["rot_2"][idx], v["rot_3"][idx]], + axis=1, + ).astype(np.float64) + + q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-12) + + # 3DGS convention is usually [w, x, y, z]. + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + + r00 = 1 - 2 * (y * y + z * z) + r01 = 2 * (x * y - w * z) + r02 = 2 * (x * z + w * y) + + r10 = 2 * (x * y + w * z) + r11 = 1 - 2 * (x * x + z * z) + r12 = 2 * (y * z - w * x) + + r20 = 2 * (x * z - w * y) + r21 = 2 * (y * z + w * x) + r22 = 1 - 2 * (x * x + y * y) + + normals = np.empty((len(idx), 3), dtype=np.float64) + + m0 = min_axis == 0 + normals[m0, 0] = r00[m0] + normals[m0, 1] = r10[m0] + normals[m0, 2] = r20[m0] + + m1 = min_axis == 1 + normals[m1, 0] = r01[m1] + normals[m1, 1] = r11[m1] + normals[m1, 2] = r21[m1] + + m2 = min_axis == 2 + normals[m2, 0] = r02[m2] + normals[m2, 1] = r12[m2] + normals[m2, 2] = r22[m2] + + # T&T trans is similarity-like. Transform directions by linear part. + A = transform_mat[:3, :3].astype(np.float64) + normals = normals @ A.T + normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12) + + return normals.astype(np.float32), scales[min_axis, np.arange(len(idx))] if False else scales + + +def estimate_gt_normals_pca(gt_points, gt_tree, gt_query_points, normal_k, batch_size): + """ + For each query point on/near GT, find K GT neighbors and estimate PCA normal. + """ + _, nn = query_tree(gt_tree, gt_query_points, k=normal_k, batch_size=batch_size) + + if nn.ndim == 1: + raise ValueError("normal_k must be > 1") + + normals = np.empty((len(gt_query_points), 3), dtype=np.float32) + + for s in range(0, len(gt_query_points), batch_size): + e = min(s + batch_size, len(gt_query_points)) + neigh = gt_points[nn[s:e]] # [B, K, 3] + centered = neigh - neigh.mean(axis=1, keepdims=True) + cov = np.einsum("bki,bkj->bij", centered, centered) / max(normal_k - 1, 1) + vals, vecs = np.linalg.eigh(cov) + n = vecs[:, :, 0] + n = n / (np.linalg.norm(n, axis=1, keepdims=True) + 1e-12) + normals[s:e] = n.astype(np.float32) + + return normals + + +def summarize_angles(angles): + return { + "normal_angular_error_mean_deg": float(np.mean(angles)), + "normal_angular_error_median_deg": float(np.median(angles)), + "normal_angular_error_q10_deg": float(np.quantile(angles, 0.10)), + "normal_angular_error_q25_deg": float(np.quantile(angles, 0.25)), + "normal_angular_error_q75_deg": float(np.quantile(angles, 0.75)), + "normal_angular_error_q90_deg": float(np.quantile(angles, 0.90)), + "normal_angular_error_iqr_deg": float(np.quantile(angles, 0.75) - np.quantile(angles, 0.25)), + } + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--method", required=True) + ap.add_argument("--scene", required=True, choices=sorted(W.SCENE_MAP.keys())) + ap.add_argument("--project-root", default="/root/autodl-tmp/SplatAtlas") + ap.add_argument("--outputs-root", default=None) + ap.add_argument("--tnt-eval-root", default=None) + ap.add_argument("--iteration", type=int, default=None) + + ap.add_argument("--mode", choices=["all", "subsample"], default="subsample") + ap.add_argument("--n-sample", type=int, default=200000) + ap.add_argument("--seed", type=int, default=0) + + ap.add_argument("--distance-multiplier", type=float, default=2.0) + ap.add_argument("--normal-k", type=int, default=30) + ap.add_argument("--max-normal-points", type=int, default=50000) + ap.add_argument("--batch-size", type=int, default=50000) + ap.add_argument("--verbose", action="store_true") + args = ap.parse_args() + + t0 = time.time() + + project_root = Path(args.project_root) + outputs_root = Path(args.outputs_root) if args.outputs_root else project_root / "outputs" + tnt_eval_root = Path(args.tnt_eval_root) if args.tnt_eval_root else project_root / "data" / "tnt_eval" + + scene_lower = args.scene.lower() + official_scene = W.SCENE_MAP[scene_lower] + tau = W.TAU_DICT[scene_lower] + near_threshold = args.distance_multiplier * tau + + ply_path = W.locate_recon_ply(outputs_root, args.method, scene_lower, args.iteration) + + scene_eval_dir = tnt_eval_root / official_scene + gt_ply_path = scene_eval_dir / f"{official_scene}.ply" + crop_path = scene_eval_dir / f"{official_scene}.json" + trans_path = scene_eval_dir / f"{official_scene}_trans.txt" + + if args.verbose: + print("=" * 80) + print("Normal angular error evaluation") + print("method:", args.method) + print("scene:", scene_lower, "->", official_scene) + print("tau:", tau) + print("near threshold:", near_threshold) + print("recon:", ply_path) + print("gt:", gt_ply_path) + + trans = W.read_transform(trans_path) + crop = W.load_crop(crop_path) + + # Load recon. + recon_raw, recon_vertex, recon_names = W.load_vertex_data(ply_path) + recon_aligned = W.apply_transform(recon_raw, trans) + recon_crop_mask = W.crop_mask_tnt(recon_aligned, crop) + recon_crop = recon_aligned[recon_crop_mask] + recon_crop_raw_idx = np.where(recon_crop_mask)[0].astype(np.int64) + + if len(recon_crop) == 0: + raise RuntimeError("Reconstruction crop is empty.") + + # Subsample reconstruction before NN and normal calculation. + eval_idx_in_crop = W.choose_eval_indices( + len(recon_crop), args.mode, args.n_sample, args.seed + ) + recon_eval = recon_crop[eval_idx_in_crop] + recon_eval_raw_idx = recon_crop_raw_idx[eval_idx_in_crop] + + # Load GT. + gt_raw, _, _ = W.load_vertex_data(gt_ply_path) + gt_crop_mask = W.crop_mask_tnt(gt_raw, crop) + gt_crop = gt_raw[gt_crop_mask] + + if len(gt_crop) == 0: + raise RuntimeError("GT crop is empty.") + + if args.verbose: + print("n_gaussians_recon:", len(recon_raw)) + print("n_recon_after_crop:", len(recon_crop)) + print("n_recon_eval:", len(recon_eval)) + print("n_gt_after_crop:", len(gt_crop)) + + # Nearest GT distance for each evaluated Gaussian center. + gt_tree = cKDTree(gt_crop) + d_r2g, nn_gt_idx = query_tree( + gt_tree, recon_eval, k=1, batch_size=max(args.batch_size, 1) + ) + + surface_mask = d_r2g < near_threshold + surface_indices_eval = np.where(surface_mask)[0].astype(np.int64) + + if args.verbose: + print("surface-near filter:", f"d < {near_threshold}") + print("n_surface_near:", len(surface_indices_eval), "/", len(recon_eval)) + if len(d_r2g): + print("distance median:", float(np.median(d_r2g))) + print("distance q10:", float(np.quantile(d_r2g, 0.10))) + print("distance q90:", float(np.quantile(d_r2g, 0.90))) + + if len(surface_indices_eval) == 0: + raise RuntimeError("No surface-near Gaussians found. Increase --distance-multiplier.") + + # Cap normal points for speed and comparable statistics. + rng = np.random.default_rng(args.seed) + if len(surface_indices_eval) > args.max_normal_points: + chosen = np.sort( + rng.choice(surface_indices_eval, size=args.max_normal_points, replace=False) + ).astype(np.int64) + else: + chosen = surface_indices_eval + + recon_normal_points = recon_eval[chosen] + recon_raw_indices_for_normals = recon_eval_raw_idx[chosen] + nearest_gt_points = gt_crop[nn_gt_idx[chosen]] + + # Gaussian normals. + gauss_normals, gauss_scales = gaussian_normals_from_vertex( + recon_vertex, + recon_names, + recon_raw_indices_for_normals, + trans, + ) + + # GT normals by local PCA. + gt_normals = estimate_gt_normals_pca( + gt_points=gt_crop, + gt_tree=gt_tree, + gt_query_points=nearest_gt_points, + normal_k=args.normal_k, + batch_size=args.batch_size, + ) + + dots = np.sum(gauss_normals * gt_normals, axis=1) + dots = np.clip(np.abs(dots), 0.0, 1.0) + angles = np.degrees(np.arccos(dots)) + + scale_min = np.min(gauss_scales, axis=1) + scale_max = np.max(gauss_scales, axis=1) + anisotropy = np.exp(scale_max - scale_min) + + result = { + "method": args.method, + "scene": scene_lower, + "official_scene": official_scene, + "eval_protocol": "auxiliary_gaussian_minor_axis_vs_tnt_gt_pca_normal", + "is_official_tnt_metric": False, + "ply_path": str(ply_path), + "gt_ply_path": str(gt_ply_path), + "crop_path": str(crop_path), + "trans_path": str(trans_path), + + "tau": float(tau), + "surface_near_threshold": float(near_threshold), + "surface_near_rule": f"d_recon_to_gt < {args.distance_multiplier} * tau", + + "mode": args.mode, + "n_sample_requested": int(args.n_sample), + "seed": int(args.seed), + "normal_k": int(args.normal_k), + "max_normal_points": int(args.max_normal_points), + + "n_gaussians_recon": int(len(recon_raw)), + "n_recon_after_crop": int(len(recon_crop)), + "n_recon_eval": int(len(recon_eval)), + "n_gt_after_crop": int(len(gt_crop)), + "n_surface_near": int(len(surface_indices_eval)), + "surface_near_ratio": float(len(surface_indices_eval) / max(len(recon_eval), 1)), + "n_normal_eval": int(len(chosen)), + + "distance_recon_to_gt_mean": float(np.mean(d_r2g)), + "distance_recon_to_gt_median": float(np.median(d_r2g)), + "distance_recon_to_gt_q10": float(np.quantile(d_r2g, 0.10)), + "distance_recon_to_gt_q90": float(np.quantile(d_r2g, 0.90)), + + **summarize_angles(angles), + + "gaussian_log_scale_min_median": float(np.median(scale_min)), + "gaussian_log_scale_max_median": float(np.median(scale_max)), + "gaussian_anisotropy_exp_scale_range_median": float(np.median(anisotropy)), + + "wall_time_seconds": float(time.time() - t0), + } + + if "opacity" in recon_names: + opacity_raw = recon_vertex["opacity"][recon_raw_indices_for_normals].astype(np.float64) + opacity_sigmoid = sigmoid(opacity_raw) + result.update({ + "opacity_raw_median": float(np.median(opacity_raw)), + "opacity_sigmoid_median": float(np.median(opacity_sigmoid)), + "opacity_sigmoid_q25": float(np.quantile(opacity_sigmoid, 0.25)), + "opacity_sigmoid_q75": float(np.quantile(opacity_sigmoid, 0.75)), + }) + + out_dir = outputs_root / "tnt_eval_normals" / f"{args.method}_{scene_lower}" + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "normal_eval.json" + with open(out_path, "w") as f: + json.dump(result, f, indent=2, sort_keys=True) + + print(json.dumps(result, indent=2, sort_keys=True)) + print(f"\n[WROTE] {out_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_tnt_normals_v3_vs_pca.py b/scripts/eval_tnt_normals_v3_vs_pca.py new file mode 100644 index 0000000000000000000000000000000000000000..6cad8981cef9edd76e303d9070248d0535a1cdd9 --- /dev/null +++ b/scripts/eval_tnt_normals_v3_vs_pca.py @@ -0,0 +1,381 @@ +#!/usr/bin/env python3 +import argparse +import json +import sys +import time +from pathlib import Path + +import numpy as np +from scipy.spatial import cKDTree + +sys.path.append(str(Path(__file__).resolve().parent)) +import eval_tnt_wrapper as W + + +def query_tree(tree, points, k=1, batch_size=50000): + ds, inds = [], [] + for s in range(0, len(points), batch_size): + e = min(s + batch_size, len(points)) + try: + d, i = tree.query(points[s:e], k=k, workers=-1) + except TypeError: + d, i = tree.query(points[s:e], k=k) + ds.append(d) + inds.append(i) + return np.concatenate(ds, axis=0), np.concatenate(inds, axis=0) + + +def pca_normals_from_neighbors(points, neighbor_indices, batch_size=50000): + """ + points: [N, 3] + neighbor_indices: [M, K] + return normals: [M, 3] + """ + if neighbor_indices.ndim != 2: + raise ValueError("neighbor_indices must be [M, K]") + + M = neighbor_indices.shape[0] + K = neighbor_indices.shape[1] + normals = np.empty((M, 3), dtype=np.float32) + + for s in range(0, M, batch_size): + e = min(s + batch_size, M) + neigh = points[neighbor_indices[s:e]] # [B, K, 3] + centered = neigh - neigh.mean(axis=1, keepdims=True) + cov = np.einsum("bki,bkj->bij", centered, centered) / max(K - 1, 1) + vals, vecs = np.linalg.eigh(cov) + n = vecs[:, :, 0] + n = n / (np.linalg.norm(n, axis=1, keepdims=True) + 1e-12) + normals[s:e] = n.astype(np.float32) + + return normals + + +def gaussian_v3_normals_from_vertex(v, names, indices, transform_mat): + required = ["scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"] + missing = [k for k in required if k not in names] + if missing: + raise ValueError(f"reconstruction PLY missing Gaussian fields: {missing}") + + idx = indices + + scales = np.stack( + [v["scale_0"][idx], v["scale_1"][idx], v["scale_2"][idx]], + axis=1, + ).astype(np.float64) + + min_axis = np.argmin(scales, axis=1) + + q = np.stack( + [v["rot_0"][idx], v["rot_1"][idx], v["rot_2"][idx], v["rot_3"][idx]], + axis=1, + ).astype(np.float64) + + q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-12) + + # 3DGS convention: quaternion [w, x, y, z] + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + + r00 = 1 - 2 * (y * y + z * z) + r01 = 2 * (x * y - w * z) + r02 = 2 * (x * z + w * y) + + r10 = 2 * (x * y + w * z) + r11 = 1 - 2 * (x * x + z * z) + r12 = 2 * (y * z - w * x) + + r20 = 2 * (x * z - w * y) + r21 = 2 * (y * z + w * x) + r22 = 1 - 2 * (x * x + y * y) + + normals = np.empty((len(idx), 3), dtype=np.float64) + + m0 = min_axis == 0 + normals[m0, 0] = r00[m0] + normals[m0, 1] = r10[m0] + normals[m0, 2] = r20[m0] + + m1 = min_axis == 1 + normals[m1, 0] = r01[m1] + normals[m1, 1] = r11[m1] + normals[m1, 2] = r21[m1] + + m2 = min_axis == 2 + normals[m2, 0] = r02[m2] + normals[m2, 1] = r12[m2] + normals[m2, 2] = r22[m2] + + # Apply linear part of T&T transform to direction. + A = transform_mat[:3, :3].astype(np.float64) + normals = normals @ A.T + normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12) + + scale_min = np.min(scales, axis=1) + scale_max = np.max(scales, axis=1) + anisotropy = np.exp(scale_max - scale_min) + + return normals.astype(np.float32), anisotropy.astype(np.float32) + + +def angular_errors(pred_normals, ref_normals): + dots = np.sum(pred_normals * ref_normals, axis=1) + dots = np.clip(np.abs(dots), 0.0, 1.0) + return np.degrees(np.arccos(dots)) + + +def stats(prefix, angles): + return { + f"{prefix}_mean_deg": float(np.mean(angles)), + f"{prefix}_median_deg": float(np.median(angles)), + f"{prefix}_q10_deg": float(np.quantile(angles, 0.10)), + f"{prefix}_q25_deg": float(np.quantile(angles, 0.25)), + f"{prefix}_q75_deg": float(np.quantile(angles, 0.75)), + f"{prefix}_q90_deg": float(np.quantile(angles, 0.90)), + f"{prefix}_iqr_deg": float(np.quantile(angles, 0.75) - np.quantile(angles, 0.25)), + } + + +def sigmoid(x): + return 1.0 / (1.0 + np.exp(-x)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--method", default="gof") + ap.add_argument("--scene", required=True, choices=sorted(W.SCENE_MAP.keys())) + ap.add_argument("--project-root", default="/root/autodl-tmp/SplatAtlas") + ap.add_argument("--outputs-root", default=None) + ap.add_argument("--tnt-eval-root", default=None) + ap.add_argument("--iteration", type=int, default=None) + + ap.add_argument("--mode", choices=["all", "subsample"], default="subsample") + ap.add_argument("--n-sample", type=int, default=200000) + ap.add_argument("--seed", type=int, default=0) + + ap.add_argument("--distance-multiplier", type=float, default=2.0) + ap.add_argument("--gt-normal-k", type=int, default=30) + ap.add_argument("--recon-pca-k", type=int, default=30) + ap.add_argument("--max-normal-points", type=int, default=50000) + ap.add_argument("--batch-size", type=int, default=50000) + ap.add_argument("--verbose", action="store_true") + args = ap.parse_args() + + t0 = time.time() + + project_root = Path(args.project_root) + outputs_root = Path(args.outputs_root) if args.outputs_root else project_root / "outputs" + tnt_eval_root = Path(args.tnt_eval_root) if args.tnt_eval_root else project_root / "data" / "tnt_eval" + + scene = args.scene.lower() + official_scene = W.SCENE_MAP[scene] + tau = W.TAU_DICT[scene] + near_threshold = args.distance_multiplier * tau + + ply_path = W.locate_recon_ply(outputs_root, args.method, scene, args.iteration) + + scene_eval_dir = tnt_eval_root / official_scene + gt_ply_path = scene_eval_dir / f"{official_scene}.ply" + crop_path = scene_eval_dir / f"{official_scene}.json" + trans_path = scene_eval_dir / f"{official_scene}_trans.txt" + + trans = W.read_transform(trans_path) + crop = W.load_crop(crop_path) + + if args.verbose: + print("=" * 80) + print("V3 normal vs Gaussian-position PCA normal") + print("method:", args.method) + print("scene:", scene, "->", official_scene) + print("tau:", tau) + print("surface-near threshold:", near_threshold) + print("recon:", ply_path) + print("gt:", gt_ply_path) + + # Load reconstruction. + recon_raw, recon_vertex, recon_names = W.load_vertex_data(ply_path) + recon_aligned = W.apply_transform(recon_raw, trans) + recon_crop_mask = W.crop_mask_tnt(recon_aligned, crop) + recon_crop = recon_aligned[recon_crop_mask] + recon_crop_raw_idx = np.where(recon_crop_mask)[0].astype(np.int64) + + if len(recon_crop) == 0: + raise RuntimeError("Reconstruction crop is empty.") + + # Choose evaluated reconstruction centers. + eval_idx_in_crop = W.choose_eval_indices( + len(recon_crop), + mode=args.mode, + n_sample=args.n_sample, + seed=args.seed, + ) + recon_eval = recon_crop[eval_idx_in_crop] + recon_eval_raw_idx = recon_crop_raw_idx[eval_idx_in_crop] + + # Load GT. + gt_raw, _, _ = W.load_vertex_data(gt_ply_path) + gt_crop_mask = W.crop_mask_tnt(gt_raw, crop) + gt_crop = gt_raw[gt_crop_mask] + + if len(gt_crop) == 0: + raise RuntimeError("GT crop is empty.") + + # Nearest GT for surface-near filter. + gt_tree = cKDTree(gt_crop) + d_r2g, nn_gt_idx = query_tree( + gt_tree, + recon_eval, + k=1, + batch_size=args.batch_size, + ) + + surface_mask = d_r2g < near_threshold + surface_eval_idx = np.where(surface_mask)[0].astype(np.int64) + + if len(surface_eval_idx) == 0: + raise RuntimeError("No surface-near Gaussians. Increase --distance-multiplier.") + + rng = np.random.default_rng(args.seed) + if len(surface_eval_idx) > args.max_normal_points: + chosen_eval_idx = np.sort( + rng.choice(surface_eval_idx, size=args.max_normal_points, replace=False) + ).astype(np.int64) + else: + chosen_eval_idx = surface_eval_idx + + chosen_points = recon_eval[chosen_eval_idx] + chosen_raw_idx = recon_eval_raw_idx[chosen_eval_idx] + chosen_nearest_gt_points = gt_crop[nn_gt_idx[chosen_eval_idx]] + + if args.verbose: + print("n_gaussians_recon:", len(recon_raw)) + print("n_recon_after_crop:", len(recon_crop)) + print("n_recon_eval:", len(recon_eval)) + print("n_gt_after_crop:", len(gt_crop)) + print("n_surface_near:", len(surface_eval_idx), "/", len(recon_eval)) + print("n_normal_eval:", len(chosen_eval_idx)) + print("distance median:", float(np.median(d_r2g))) + print("distance q10:", float(np.quantile(d_r2g, 0.10))) + print("distance q90:", float(np.quantile(d_r2g, 0.90))) + + # Reference: GT local PCA normals at nearest GT points. + if args.verbose: + print("[GT PCA normals] k =", args.gt_normal_k) + + _, gt_neighbor_idx = query_tree( + gt_tree, + chosen_nearest_gt_points, + k=args.gt_normal_k, + batch_size=args.batch_size, + ) + gt_pca_normals = pca_normals_from_neighbors( + gt_crop, + gt_neighbor_idx, + batch_size=args.batch_size, + ) + + # Method 1: V3 / Gaussian minor-axis normal. + if args.verbose: + print("[V3 normals] from Gaussian scale + rotation") + + v3_normals, anisotropy = gaussian_v3_normals_from_vertex( + recon_vertex, + recon_names, + chosen_raw_idx, + trans, + ) + + v3_angles = angular_errors(v3_normals, gt_pca_normals) + + # Method 2: Gaussian-position local PCA normal. + if args.verbose: + print("[Gaussian-position PCA normals] k =", args.recon_pca_k) + + recon_tree = cKDTree(recon_crop) + _, recon_neighbor_idx = query_tree( + recon_tree, + chosen_points, + k=args.recon_pca_k, + batch_size=args.batch_size, + ) + recon_pca_normals = pca_normals_from_neighbors( + recon_crop, + recon_neighbor_idx, + batch_size=args.batch_size, + ) + + recon_pca_angles = angular_errors(recon_pca_normals, gt_pca_normals) + + result = { + "method": args.method, + "scene": scene, + "official_scene": official_scene, + "eval_protocol": "gof_v3_minor_axis_vs_gaussian_position_pca_normal", + "is_official_tnt_metric": False, + + "ply_path": str(ply_path), + "gt_ply_path": str(gt_ply_path), + "crop_path": str(crop_path), + "trans_path": str(trans_path), + + "tau": float(tau), + "surface_near_threshold": float(near_threshold), + "surface_near_rule": f"d_recon_to_gt < {args.distance_multiplier} * tau", + + "mode": args.mode, + "n_sample_requested": int(args.n_sample), + "seed": int(args.seed), + "gt_normal_k": int(args.gt_normal_k), + "recon_pca_k": int(args.recon_pca_k), + "max_normal_points": int(args.max_normal_points), + + "n_gaussians_recon": int(len(recon_raw)), + "n_recon_after_crop": int(len(recon_crop)), + "n_recon_eval": int(len(recon_eval)), + "n_gt_after_crop": int(len(gt_crop)), + "n_surface_near": int(len(surface_eval_idx)), + "surface_near_ratio": float(len(surface_eval_idx) / max(len(recon_eval), 1)), + "n_normal_eval": int(len(chosen_eval_idx)), + + "distance_recon_to_gt_mean": float(np.mean(d_r2g)), + "distance_recon_to_gt_median": float(np.median(d_r2g)), + "distance_recon_to_gt_q10": float(np.quantile(d_r2g, 0.10)), + "distance_recon_to_gt_q90": float(np.quantile(d_r2g, 0.90)), + + **stats("v3_normal_error", v3_angles), + **stats("pca_normal_error", recon_pca_angles), + + "median_delta_v3_minus_pca_deg": float(np.median(v3_angles) - np.median(recon_pca_angles)), + "mean_delta_v3_minus_pca_deg": float(np.mean(v3_angles) - np.mean(recon_pca_angles)), + "pca_better_than_v3_by_median": bool(np.median(recon_pca_angles) < np.median(v3_angles)), + "pca_within_3deg_of_v3_median": bool(np.median(recon_pca_angles) <= np.median(v3_angles) + 3.0), + "pca_within_5deg_of_v3_median": bool(np.median(recon_pca_angles) <= np.median(v3_angles) + 5.0), + + "gaussian_anisotropy_exp_scale_range_median": float(np.median(anisotropy)), + "gaussian_anisotropy_exp_scale_range_q25": float(np.quantile(anisotropy, 0.25)), + "gaussian_anisotropy_exp_scale_range_q75": float(np.quantile(anisotropy, 0.75)), + + "wall_time_seconds": float(time.time() - t0), + } + + if "opacity" in recon_names: + opacity_raw = recon_vertex["opacity"][chosen_raw_idx].astype(np.float64) + opacity_sigmoid = sigmoid(opacity_raw) + result.update({ + "opacity_sigmoid_median": float(np.median(opacity_sigmoid)), + "opacity_sigmoid_q25": float(np.quantile(opacity_sigmoid, 0.25)), + "opacity_sigmoid_q75": float(np.quantile(opacity_sigmoid, 0.75)), + }) + + out_dir = outputs_root / "tnt_eval_normals_v3_vs_pca" / f"{args.method}_{scene}" + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "v3_vs_pca_normal_eval.json" + + with open(out_path, "w") as f: + json.dump(result, f, indent=2, sort_keys=True) + + print(json.dumps(result, indent=2, sort_keys=True)) + print(f"\n[WROTE] {out_path}") + + +if __name__ == "__main__": + main() diff --git a/scripts/eval_tnt_wrapper.py b/scripts/eval_tnt_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..7daae712e9e51cfa7d7e79f30ccd7759a9cafd69 --- /dev/null +++ b/scripts/eval_tnt_wrapper.py @@ -0,0 +1,528 @@ +#!/usr/bin/env python3 +import argparse +import json +import math +import sys +import time +from pathlib import Path + +import numpy as np +from plyfile import PlyData +from scipy.spatial import cKDTree + + +SCENE_MAP = { + "barn": "Barn", + "caterpillar": "Caterpillar", + "truck": "Truck", + "church": "Church", + "courthouse": "Courthouse", + "ignatius": "Ignatius", + "meetingroom": "Meetingroom", +} + +TAU_DICT = { + "barn": 0.01, + "caterpillar": 0.005, + "truck": 0.005, + "church": 0.01, + "courthouse": 0.01, + "ignatius": 0.003, + "meetingroom": 0.01, +} + +ITER_DICT = { + "gaussian_surfel": 15000, + "hogs": 50000, +} + +DEFAULT_ITER = 30000 + + +def vlog(enabled, *args): + if enabled: + print(*args, flush=True) + + +def as_float(x): + if isinstance(x, (np.floating, np.integer)): + return x.item() + if isinstance(x, np.ndarray): + return x.tolist() + return x + + +def locate_recon_ply(outputs_root: Path, method: str, scene: str, iteration=None) -> Path: + scene = scene.lower() + + if iteration is None: + iteration = ITER_DICT.get(method, DEFAULT_ITER) + + candidates = [] + + # Important special case from user's training history. + if method == "vanilla_3dgs": + candidates.append( + outputs_root + / f"{method}_{scene}_bak" + / "point_cloud" + / f"iteration_{iteration}" + / "point_cloud.ply" + ) + + candidates.append( + outputs_root + / f"{method}_{scene}" + / "point_cloud" + / f"iteration_{iteration}" + / "point_cloud.ply" + ) + + # Fallback: if requested iteration is absent, search any iteration folder. + for base in [outputs_root / f"{method}_{scene}_bak", outputs_root / f"{method}_{scene}"]: + if base.exists(): + candidates.extend(sorted(base.glob("point_cloud/iteration_*/point_cloud.ply"))) + + for p in candidates: + if p.exists(): + return p + + msg = "Could not find reconstruction PLY. Tried:\n" + "\n".join(str(p) for p in candidates) + raise FileNotFoundError(msg) + + +def load_vertex_data(ply_path: Path): + ply = PlyData.read(str(ply_path)) + if "vertex" not in ply: + raise ValueError(f"No vertex element in PLY: {ply_path}") + v = ply["vertex"].data + names = v.dtype.names + for k in ["x", "y", "z"]: + if k not in names: + raise ValueError(f"PLY missing vertex property '{k}': {ply_path}") + xyz = np.stack([v["x"], v["y"], v["z"]], axis=1).astype(np.float32, copy=False) + return xyz, v, names + + +def load_xyz_only(ply_path: Path): + xyz, _, _ = load_vertex_data(ply_path) + return xyz + + +def read_transform(path: Path): + mat = np.loadtxt(path).astype(np.float64) + if mat.shape != (4, 4): + raise ValueError(f"Transform must be 4x4, got {mat.shape}: {path}") + return mat + + +def apply_transform(xyz: np.ndarray, mat: np.ndarray): + # Homogeneous column convention: x_gt = M @ [x, y, z, 1]. + out = xyz.astype(np.float64) @ mat[:3, :3].T + mat[:3, 3][None, :] + return out.astype(np.float32) + + +def axis_to_index(axis): + if isinstance(axis, int): + return axis + s = str(axis).lower() + if s in ["0", "x", "axis_x"]: + return 0 + if s in ["1", "y", "axis_y"]: + return 1 + if s in ["2", "z", "axis_z"]: + return 2 + raise ValueError(f"Unknown orthogonal_axis: {axis}") + + +def points_in_polygon(points2: np.ndarray, polygon2: np.ndarray): + # Try matplotlib's optimized path first. Fallback to vectorized ray casting. + try: + from matplotlib.path import Path as MplPath + return MplPath(polygon2).contains_points(points2) + except Exception: + x = points2[:, 0] + y = points2[:, 1] + poly = polygon2 + inside = np.zeros(points2.shape[0], dtype=bool) + j = len(poly) - 1 + for i in range(len(poly)): + xi, yi = poly[i] + xj, yj = poly[j] + denom = (yj - yi) + if abs(denom) < 1e-30: + denom = 1e-30 + intersect = ((yi > y) != (yj > y)) & ( + x < (xj - xi) * (y - yi) / denom + xi + ) + inside ^= intersect + j = i + return inside + + +def crop_mask_tnt(xyz: np.ndarray, crop_json: dict): + axis = axis_to_index(crop_json["orthogonal_axis"]) + axis_min = float(crop_json["axis_min"]) + axis_max = float(crop_json["axis_max"]) + + axis_mask = (xyz[:, axis] >= axis_min) & (xyz[:, axis] <= axis_max) + + poly = np.asarray(crop_json["bounding_polygon"], dtype=np.float64) + other_axes = [i for i in [0, 1, 2] if i != axis] + + if poly.ndim != 2: + raise ValueError(f"Invalid bounding_polygon shape: {poly.shape}") + + if poly.shape[1] == 3: + poly2 = poly[:, other_axes] + elif poly.shape[1] == 2: + poly2 = poly + else: + raise ValueError(f"Invalid bounding_polygon shape: {poly.shape}") + + idx = np.where(axis_mask)[0] + final_mask = np.zeros(xyz.shape[0], dtype=bool) + + if idx.size == 0: + return final_mask + + pts2 = xyz[idx][:, other_axes] + inside = points_in_polygon(pts2, poly2) + final_mask[idx] = inside + return final_mask + + +def load_crop(path: Path): + with open(path, "r") as f: + return json.load(f) + + +def build_tree(points: np.ndarray, name: str, verbose=False): + if len(points) == 0: + raise ValueError(f"Cannot build KDTree for empty point set: {name}") + vlog(verbose, f"[KDTree] building {name}: {len(points):,} points") + return cKDTree(points) + + +def query_tree(tree, points, batch_size=250000): + n = len(points) + dists = np.empty(n, dtype=np.float32) + inds = np.empty(n, dtype=np.int64) + + for s in range(0, n, batch_size): + e = min(s + batch_size, n) + try: + d, idx = tree.query(points[s:e], k=1, workers=-1) + except TypeError: + d, idx = tree.query(points[s:e], k=1) + dists[s:e] = d.astype(np.float32) + inds[s:e] = idx.astype(np.int64) + + return dists, inds + + +def compute_metrics(recon: np.ndarray, gt: np.ndarray, tau: float, batch_size=250000, verbose=False): + gt_tree = build_tree(gt, "GT", verbose) + d_recon_to_gt, _ = query_tree(gt_tree, recon, batch_size=batch_size) + precision = float(np.mean(d_recon_to_gt < tau)) + chamfer_r2g = float(np.mean(d_recon_to_gt)) + + recon_tree = build_tree(recon, "reconstruction", verbose) + d_gt_to_recon, _ = query_tree(recon_tree, gt, batch_size=batch_size) + recall = float(np.mean(d_gt_to_recon < tau)) + chamfer_g2r = float(np.mean(d_gt_to_recon)) + + if precision + recall > 0: + f_score = float(2.0 * precision * recall / (precision + recall)) + else: + f_score = 0.0 + + # User requested "mean of bidirectional"; keep sum too for transparency. + chamfer_mean = float(0.5 * (chamfer_r2g + chamfer_g2r)) + chamfer_sum = float(chamfer_r2g + chamfer_g2r) + + return { + "precision": precision, + "recall": recall, + "f_score": f_score, + "chamfer_distance": chamfer_mean, + "chamfer_distance_sum": chamfer_sum, + "chamfer_recon_to_gt": chamfer_r2g, + "chamfer_gt_to_recon": chamfer_g2r, + } + + +def choose_eval_indices(n: int, mode: str, n_sample: int, seed: int): + if mode == "all": + return np.arange(n, dtype=np.int64) + + if mode == "subsample": + k = min(n, n_sample) + rng = np.random.default_rng(seed) + return np.sort(rng.choice(n, size=k, replace=False)).astype(np.int64) + + raise ValueError(f"Unsupported mode: {mode}") + + +def get_normals_from_gt_vertex(v, names): + candidates = [ + ("nx", "ny", "nz"), + ("normal_x", "normal_y", "normal_z"), + ] + for a, b, c in candidates: + if a in names and b in names and c in names: + normals = np.stack([v[a], v[b], v[c]], axis=1).astype(np.float32) + denom = np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12 + return normals / denom + return None + + +def gaussian_normals_from_vertex(v, names, indices, transform_mat): + required = ["scale_0", "scale_1", "scale_2", "rot_0", "rot_1", "rot_2", "rot_3"] + if not all(k in names for k in required): + return None + + idx = indices + scales = np.stack([v["scale_0"][idx], v["scale_1"][idx], v["scale_2"][idx]], axis=1) + min_axis = np.argmin(scales, axis=1) + + q = np.stack([v["rot_0"][idx], v["rot_1"][idx], v["rot_2"][idx], v["rot_3"][idx]], axis=1).astype(np.float64) + q = q / (np.linalg.norm(q, axis=1, keepdims=True) + 1e-12) + + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + + # Rotation matrix for quaternion [w, x, y, z]. + r00 = 1 - 2 * (y * y + z * z) + r01 = 2 * (x * y - w * z) + r02 = 2 * (x * z + w * y) + + r10 = 2 * (x * y + w * z) + r11 = 1 - 2 * (x * x + z * z) + r12 = 2 * (y * z - w * x) + + r20 = 2 * (x * z - w * y) + r21 = 2 * (y * z + w * x) + r22 = 1 - 2 * (x * x + y * y) + + normals = np.empty((len(idx), 3), dtype=np.float64) + + mask0 = min_axis == 0 + normals[mask0, 0] = r00[mask0] + normals[mask0, 1] = r10[mask0] + normals[mask0, 2] = r20[mask0] + + mask1 = min_axis == 1 + normals[mask1, 0] = r01[mask1] + normals[mask1, 1] = r11[mask1] + normals[mask1, 2] = r21[mask1] + + mask2 = min_axis == 2 + normals[mask2, 0] = r02[mask2] + normals[mask2, 1] = r12[mask2] + normals[mask2, 2] = r22[mask2] + + # For near-similarity transforms this is equivalent up to scale. + A = transform_mat[:3, :3].astype(np.float64) + normals = normals @ A.T + normals = normals / (np.linalg.norm(normals, axis=1, keepdims=True) + 1e-12) + return normals.astype(np.float32) + + +def compute_normal_error( + recon_eval_points, + recon_eval_raw_indices, + recon_vertex, + recon_names, + gt_crop_points, + gt_crop_normals, + transform_mat, + batch_size, +): + if gt_crop_normals is None: + return {"normal_error_available": False, "normal_error_reason": "GT PLY has no normal fields"} + + recon_normals = gaussian_normals_from_vertex( + recon_vertex, recon_names, recon_eval_raw_indices, transform_mat + ) + if recon_normals is None: + return { + "normal_error_available": False, + "normal_error_reason": "reconstruction PLY has no scale_*/rot_* fields", + } + + gt_tree = cKDTree(gt_crop_points) + _, nn_idx = query_tree(gt_tree, recon_eval_points, batch_size=batch_size) + nn_gt_normals = gt_crop_normals[nn_idx] + + dots = np.sum(recon_normals * nn_gt_normals, axis=1) + dots = np.clip(np.abs(dots), 0.0, 1.0) + angles = np.degrees(np.arccos(dots)) + + return { + "normal_error_available": True, + "normal_angular_error_median_deg": float(np.median(angles)), + "normal_angular_error_q25_deg": float(np.quantile(angles, 0.25)), + "normal_angular_error_q75_deg": float(np.quantile(angles, 0.75)), + "normal_angular_error_iqr_deg": float(np.quantile(angles, 0.75) - np.quantile(angles, 0.25)), + "normal_error_n_points": int(len(angles)), + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--method", required=True) + parser.add_argument("--scene", required=True, choices=sorted(SCENE_MAP.keys())) + parser.add_argument("--project-root", default="/root/autodl-tmp/SplatAtlas") + parser.add_argument("--outputs-root", default=None) + parser.add_argument("--tnt-eval-root", default=None) + parser.add_argument("--iteration", type=int, default=None) + parser.add_argument("--mode", choices=["all", "subsample"], default="subsample") + parser.add_argument("--n-sample", type=int, default=200000) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--batch-size", type=int, default=250000) + parser.add_argument("--normal-error", action="store_true") + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args() + + t0 = time.time() + + project_root = Path(args.project_root) + outputs_root = Path(args.outputs_root) if args.outputs_root else project_root / "outputs" + tnt_eval_root = Path(args.tnt_eval_root) if args.tnt_eval_root else project_root / "data" / "tnt_eval" + + scene_lower = args.scene.lower() + official_scene = SCENE_MAP[scene_lower] + tau = TAU_DICT[scene_lower] + + ply_path = locate_recon_ply(outputs_root, args.method, scene_lower, args.iteration) + + scene_eval_dir = tnt_eval_root / official_scene + gt_ply_path = scene_eval_dir / f"{official_scene}.ply" + crop_path = scene_eval_dir / f"{official_scene}.json" + trans_path = scene_eval_dir / f"{official_scene}_trans.txt" + + for p in [gt_ply_path, crop_path, trans_path]: + if not p.exists(): + raise FileNotFoundError(f"Missing required eval file: {p}") + + vlog(args.verbose, "=" * 80) + vlog(args.verbose, f"method: {args.method}") + vlog(args.verbose, f"scene: {scene_lower} -> {official_scene}") + vlog(args.verbose, f"tau: {tau}") + vlog(args.verbose, f"recon ply: {ply_path}") + vlog(args.verbose, f"gt ply: {gt_ply_path}") + vlog(args.verbose, f"crop json: {crop_path}") + vlog(args.verbose, f"trans: {trans_path}") + + trans = read_transform(trans_path) + crop_json = load_crop(crop_path) + + vlog(args.verbose, "\n[Load] reconstruction PLY") + recon_xyz_raw, recon_vertex, recon_names = load_vertex_data(ply_path) + n_gaussians_recon = int(len(recon_xyz_raw)) + vlog(args.verbose, f"n_gaussians_recon: {n_gaussians_recon:,}") + + vlog(args.verbose, "\n[Transform] applying T&T trans matrix") + vlog(args.verbose, trans) + recon_xyz_aligned = apply_transform(recon_xyz_raw, trans) + + vlog(args.verbose, "\n[Crop] reconstruction") + recon_crop_mask = crop_mask_tnt(recon_xyz_aligned, crop_json) + recon_crop_indices_raw = np.where(recon_crop_mask)[0].astype(np.int64) + recon_crop = recon_xyz_aligned[recon_crop_mask] + vlog(args.verbose, f"recon before crop: {len(recon_xyz_aligned):,}") + vlog(args.verbose, f"recon after crop: {len(recon_crop):,}") + + if len(recon_crop) == 0: + raise RuntimeError("Reconstruction crop is empty. Likely coordinate-system/alignment mismatch.") + + vlog(args.verbose, "\n[Load] GT PLY") + gt_xyz_raw, gt_vertex, gt_names = load_vertex_data(gt_ply_path) + vlog(args.verbose, f"GT raw points: {len(gt_xyz_raw):,}") + + vlog(args.verbose, "\n[Crop] GT") + gt_crop_mask = crop_mask_tnt(gt_xyz_raw, crop_json) + gt_crop = gt_xyz_raw[gt_crop_mask] + vlog(args.verbose, f"GT after crop: {len(gt_crop):,}") + + if len(gt_crop) == 0: + raise RuntimeError("GT crop is empty. Crop parser is likely wrong.") + + eval_idx_in_crop = choose_eval_indices( + len(recon_crop), mode=args.mode, n_sample=args.n_sample, seed=args.seed + ) + recon_eval = recon_crop[eval_idx_in_crop] + recon_eval_raw_indices = recon_crop_indices_raw[eval_idx_in_crop] + + vlog(args.verbose, "\n[Eval set]") + vlog(args.verbose, f"mode: {args.mode}") + vlog(args.verbose, f"recon eval points: {len(recon_eval):,}") + vlog(args.verbose, f"GT eval points: {len(gt_crop):,}") + + metrics = compute_metrics( + recon_eval, + gt_crop, + tau=tau, + batch_size=args.batch_size, + verbose=args.verbose, + ) + + normal_metrics = {} + if args.normal_error: + vlog(args.verbose, "\n[Normal error]") + gt_normals_raw = get_normals_from_gt_vertex(gt_vertex, gt_names) + gt_normals_crop = gt_normals_raw[gt_crop_mask] if gt_normals_raw is not None else None + normal_metrics = compute_normal_error( + recon_eval_points=recon_eval, + recon_eval_raw_indices=recon_eval_raw_indices, + recon_vertex=recon_vertex, + recon_names=recon_names, + gt_crop_points=gt_crop, + gt_crop_normals=gt_normals_crop, + transform_mat=trans, + batch_size=args.batch_size, + ) + vlog(args.verbose, normal_metrics) + + wall_time = time.time() - t0 + + result = { + "method": args.method, + "scene": scene_lower, + "official_scene": official_scene, + "eval_protocol": "tnt_point_cloud_direct_gaussian_centers", + "mode": args.mode, + "n_sample_requested": int(args.n_sample), + "seed": int(args.seed), + "ply_path": str(ply_path), + "gt_ply_path": str(gt_ply_path), + "crop_path": str(crop_path), + "trans_path": str(trans_path), + "n_gaussians_recon": n_gaussians_recon, + "n_points_after_crop": int(len(recon_crop)), + "n_points_eval_recon": int(len(recon_eval)), + "n_points_gt_raw": int(len(gt_xyz_raw)), + "n_points_gt": int(len(gt_crop)), + "tau": float(tau), + **metrics, + **normal_metrics, + "wall_time_seconds": float(wall_time), + } + + out_dir = outputs_root / "tnt_eval" / f"{args.method}_{scene_lower}" + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "eval.json" + + with open(out_path, "w") as f: + json.dump(result, f, indent=2, sort_keys=True) + + print(json.dumps(result, indent=2, sort_keys=True)) + print(f"\n[WROTE] {out_path}") + + +if __name__ == "__main__": + try: + main() + except Exception as e: + print(f"[ERROR] {type(e).__name__}: {e}", file=sys.stderr) + raise diff --git a/scripts/extract_dynamics.py b/scripts/extract_dynamics.py new file mode 100644 index 0000000000000000000000000000000000000000..4589fc66a5b00fc4252bd4bbc499c5eca61672dc --- /dev/null +++ b/scripts/extract_dynamics.py @@ -0,0 +1,159 @@ +import os +import glob +import pandas as pd +import matplotlib.pyplot as plt +import seaborn as sns +from tensorboard.backend.event_processing import event_accumulator + +# === 论文级图表样式设置 === +sns.set_theme(style="whitegrid", palette="muted") +plt.rcParams.update({ + 'font.size': 12, + 'axes.titlesize': 14, + 'axes.labelsize': 12, + 'figure.autolayout': True, + 'figure.dpi': 300 +}) + +def extract_tb_data(log_dir): + """从 TensorBoard 目录中提取所有 Scalar 数据为 Pandas DataFrame""" + event_files = glob.glob(os.path.join(log_dir, 'events.out.tfevents.*')) + if not event_files: + return None + + # size_guidance=0 保证我们不会截断任何一个 step 的数据 + ea = event_accumulator.EventAccumulator(log_dir, size_guidance={event_accumulator.SCALARS: 0}) + ea.Reload() + + extracted_data = {} + if 'scalars' not in ea.Tags(): + return None + + for tag in ea.Tags()['scalars']: + events = ea.Scalars(tag) + # 将每个指标转换为 DataFrame + extracted_data[tag] = pd.DataFrame( + [(e.step, e.value) for e in events], + columns=['step', tag] + ).set_index('step') + + if not extracted_data: + return None + + # 以 step 为基准，横向拼接所有监控指标 + df = pd.concat(extracted_data.values(), axis=1).reset_index() + return df + +def main(): + outputs_dir = os.path.join(os.path.dirname(__file__), "..", "outputs") + plots_dir = os.path.join(os.path.dirname(__file__), "..", "plots") + os.makedirs(plots_dir, exist_ok=True) + + all_data = [] + + print("🔍 [Harvester] 正在扫描 TensorBoard 日志...") + for folder in os.listdir(outputs_dir): + log_dir = os.path.join(outputs_dir, folder) + # 忽略非目录和病理学隔离区（可选，这里为了看全盘先全扫） + if not os.path.isdir(log_dir) or folder == "pathology_archive": + continue + + df = extract_tb_data(log_dir) + if df is not None: + # 假设文件夹命名规范为 method_scene + parts = folder.split('_', 1) + method = parts[0] if len(parts) > 1 else folder + + df['method'] = method + df['experiment'] = folder + all_data.append(df) + print(f" ✅ 成功提取: {folder} ({len(df)} 条时间序列记录)") + + if not all_data: + print("⚠️ 未找到任何有效的 TensorBoard 数据！") + return + + master_df = pd.concat(all_data, ignore_index=True) + + # 动态匹配我们埋下的探针列名 + time_col = 'hardware/cumulative_time_sec' if 'hardware/cumulative_time_sec' in master_df.columns else 'step' + n_col = 'geometry/N_gaussians' + loss_col = 'train/loss' + cos_col = 'dynamics_decoupled/cos_similarity' + gamma_col = 'pathology/gamma_median' + + print(f"\n📊 [Harvester] 正在使用 Matplotlib 绘制动力学大图，目标目录: {plots_dir}/") + + # ========================================================= + # 图表 1: 边际收益崩塌与致密化 S 曲线 (The Marginal Return Curve) + # ========================================================= + if n_col in master_df.columns and loss_col in master_df.columns: + plt.figure(figsize=(10, 6)) + ax1 = plt.gca() + ax2 = ax1.twinx() + + colors = sns.color_palette("husl", len(master_df['experiment'].unique())) + + for idx, exp in enumerate(master_df['experiment'].unique()): + exp_df = master_df[master_df['experiment'] == exp].dropna(subset=[time_col, n_col, loss_col]) + if exp_df.empty: continue + + color = colors[idx] + # 绘制 N 数量的 S 曲线 (实线) + ax1.plot(exp_df[time_col], exp_df[n_col], linestyle='-', linewidth=2.5, color=color, label=f"{exp} (N)") + # 绘制 Loss 下降曲线 (虚线) + ax2.plot(exp_df[time_col], exp_df[loss_col], linestyle='--', linewidth=2, color=color, alpha=0.6) + + xlabel = 'Cumulative GPU Time (Seconds)' if time_col != 'step' else 'Training Steps' + ax1.set_xlabel(xlabel, fontweight='bold') + ax1.set_ylabel('Number of Gaussians (N)', color='black', fontweight='bold') + ax2.set_ylabel('Training Loss (L1 + SSIM)', color='gray', fontweight='bold') + + plt.title('Densification Dynamics & Loss Convergence', pad=20, fontweight='bold') + ax1.legend(loc='upper left', bbox_to_anchor=(1.15, 1)) + + out_path1 = os.path.join(plots_dir, 'fig1_marginal_return.png') + plt.savefig(out_path1, bbox_inches='tight') + plt.close() + print(f" 📸 论文大图 1 已生成: {out_path1}") + + # ========================================================= + # 图表 2: 光度-几何代偿动力学散点图 (The Compensation Scatter) + # ========================================================= + if cos_col in master_df.columns and gamma_col in master_df.columns: + plot_df = master_df.dropna(subset=[cos_col, gamma_col]).copy() + if not plot_df.empty: + plt.figure(figsize=(9, 7)) + + # 使用散点图，横坐标是算子夹角，纵坐标是畸变率，点的大小代表训练进度(Step) + scatter = sns.scatterplot( + data=plot_df, + x=cos_col, + y=gamma_col, + hue='method', + size='step', + sizes=(20, 200), + alpha=0.8, + palette='Set1', + edgecolor='black' + ) + + # 增加物理意义的警戒线 + plt.axvline(x=0, color='gray', linestyle='--', alpha=0.5, label='Orthogonal Gradients (0)') + plt.axhline(y=18, color='red', linestyle=':', linewidth=2, alpha=0.7, label=r'Pathology Threshold ($\gamma=18$)') + + plt.xlabel(r'Gradient Cosine Similarity $\cos(\theta)$ (Target vs Parasitic)', fontweight='bold') + plt.ylabel(r'Anisotropy Distortion $\gamma$ (Max/Min Scale)', fontweight='bold') + plt.title('Photometric-Geometric Compensation Phase Space', pad=20, fontweight='bold') + + plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left') + + out_path2 = os.path.join(plots_dir, 'fig2_compensation_scatter.png') + plt.savefig(out_path2, bbox_inches='tight') + plt.close() + print(f" 📸 论文大图 2 已生成: {out_path2}") + + print("✅ [Harvester] 您的核心科研图表已就绪！") + +if __name__ == "__main__": + main() diff --git a/scripts/fix_bonsai_defaults.py b/scripts/fix_bonsai_defaults.py new file mode 100644 index 0000000000000000000000000000000000000000..77b1932f89c7f7853c23a2f64b4f209d9713c425 --- /dev/null +++ b/scripts/fix_bonsai_defaults.py @@ -0,0 +1,52 @@ +import os, subprocess + +SPLATATLAS = "/root/autodl-tmp/SplatAtlas" +BASE = "/root/autodl-tmp/SplatAtlas/outputs" +SRC = "/root/autodl-tmp/dataset/360/bonsai" +RES = 2 +SCENE = "bonsai" + +ENV_MAP = { + "lightgaussian": ("/root/autodl-tmp/envs/lightgaussian_official", "/root/autodl-tmp/LightGaussian_official"), + "3dgsmcmc": ("/root/autodl-tmp/envs/3dgs_mcmc", "/root/autodl-tmp/3dgs-mcmc") +} + +def get_py_and_env(method): + env_path, extra_py = ENV_MAP[method] + python_bin = f"{env_path}/bin/python" + env = os.environ.copy() + env["PYTHONPATH"] = f"{SPLATATLAS}:{extra_py}:{env.get('PYTHONPATH', '')}" + env["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" # 防止算 LPIPS 爆显存 + return python_bin, env + +targets = ["lightgaussian", "3dgsmcmc"] + +for method in targets: + m_dir = f"{BASE}/{method}_{SCENE}" + metrics_f = f"{m_dir}/metrics_test_iter30000.json" + ply_f = f"{m_dir}/point_cloud/iteration_30000/point_cloud.ply" + + if not os.path.exists(ply_f): + print(f"[-] SKIPPING {method}_{SCENE}: PLY 未找到，说明训练没跑完。") + continue + + print(f"\n[*] FIXING: {method}_{SCENE}") + python_bin, custom_env = get_py_and_env(method) + + # 1. 强制渲染 + render_cmd = f"{python_bin} {SPLATATLAS}/scripts/main_render.py --method {method} --source_path {SRC} --model_path {m_dir} --iteration 30000 --resolution {RES}" + subprocess.run(render_cmd, shell=True, env=custom_env) + + # 2. 评估 + renders = f"{m_dir}/renders_test_30000" + gt = f"{m_dir}/gt_test_30000" + depths = f"{m_dir}/depths_test_30000" + + eval_cmd = ( + f"{python_bin} {SPLATATLAS}/ufd_evalkit/run_eval.py " + f"--method {method} --scene {SCENE} --render_dir {renders} " + f"--gt_dir {gt} --ply_path {ply_f} --output_json {metrics_f} " + f"--colmap_dir {SRC} --depth_dir {depths}" + ) + subprocess.run(eval_cmd, shell=True, env=custom_env) + diff --git a/scripts/fix_colmap_format.py b/scripts/fix_colmap_format.py new file mode 100644 index 0000000000000000000000000000000000000000..4d4a977a5572acf546252edae81a10de3c80ae76 --- /dev/null +++ b/scripts/fix_colmap_format.py @@ -0,0 +1,56 @@ +import sys +import os + +# 动态将 taming-3dgs 路径加入，利用其内置的 colmap_loader +sys.path.append("/root/autodl-tmp/taming-3dgs") +try: + from scene.colmap_loader import read_extrinsics_binary, read_intrinsics_binary, qvec2rotmat + from scene.colmap_loader import read_points3D_binary +except ImportError: + print("❌ 找不到 taming-3dgs 路径，请确认目录存在。") + sys.exit(1) + +def convert(): + base_path = "/root/autodl-tmp/dataset/synthetic_nerf/Synthetic_NeRF/Lego/sparse/0" + print(f"🔄 正在手动解析二进制并生成文本文件: {base_path}") + + # 1. 处理相机的内参 (Intrinsics) + try: + intrinsics = read_intrinsics_binary(os.path.join(base_path, "cameras.bin")) + with open(os.path.join(base_path, "cameras.txt"), "w") as f: + for i in intrinsics: + # 格式: CAMERA_ID MODEL WIDTH HEIGHT PARAMS[] + params = " ".join([str(p) for p in intrinsics[i].params]) + f.write(f"{i} {intrinsics[i].model} {intrinsics[i].width} {intrinsics[i].height} {params}\n") + print(" - cameras.txt 生成成功") + except Exception as e: print(f" - cameras 转换失败: {e}") + + # 2. 处理外参和图像 (Extrinsics/Images) + try: + extrinsics = read_extrinsics_binary(os.path.join(base_path, "images.bin")) + with open(os.path.join(base_path, "images.txt"), "w") as f: + for i in extrinsics: + e = extrinsics[i] + q = " ".join([str(x) for x in e.qvec]) + t = " ".join([str(x) for x in e.tvec]) + # 第一行: IMAGE_ID QW QX QY QZ TX TY TZ CAMERA_ID NAME + f.write(f"{i} {q} {t} {e.camera_id} {e.name}\n") + # 第二行: POINTS2D (Blender 数据通常是空的或占位符) + f.write("\n") + print(" - images.txt 生成成功") + except Exception as e: print(f" - images 转换失败: {e}") + + # 3. 处理点云 (Points3D) + try: + pts3d = read_points3D_binary(os.path.join(base_path, "points3D.bin")) + with open(os.path.join(base_path, "points3D.txt"), "w") as f: + for i in pts3d: + p = pts3d[i] + xyz = " ".join([str(x) for x in p.xyz]) + rgb = " ".join([str(x) for x in p.rgb]) + f.write(f"{i} {xyz} {rgb} {p.error} 0\n") + print(" - points3D.txt 生成成功") + except Exception as e: print(f" - points3D 转换失败: {e}") + +if __name__ == "__main__": + convert() diff --git a/scripts/fix_missing_evals.py b/scripts/fix_missing_evals.py new file mode 100644 index 0000000000000000000000000000000000000000..4d57513bd8282f6b9cabf87799c1a15daedd0b4d --- /dev/null +++ b/scripts/fix_missing_evals.py @@ -0,0 +1,47 @@ +import os, subprocess + +SPLATATLAS = "/root/autodl-tmp/SplatAtlas" +BASE = "/root/autodl-tmp/SplatAtlas/outputs" +PYTHON = "python" # 默认先用 base 环境 + +targets = [ + ("vanilla_3dgs", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2), + ("analyticsplatting", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2), + ("analyticsplatting", "Lego", "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/Lego", 1), + ("erankgs", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2), + ("pgsr", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2) +] +seeds = ["", "_seed1", "_seed2"] + +for method, scene, src, res in targets: + for suffix in seeds: + m_dir = f"{BASE}/{method}{suffix}_{scene}" + metrics_f = f"{m_dir}/metrics_test_iter30000.json" + ply_f = f"{m_dir}/point_cloud/iteration_30000/point_cloud.ply" + + if os.path.exists(metrics_f): + continue + + if not os.path.exists(ply_f): + print(f"[-] SKIPPING {method}{suffix}_{scene}: PLY not found at {ply_f}") + continue + + print(f"[*] FIXING: {method}{suffix}_{scene}") + + # 1. 渲染 (强制覆盖) + render_cmd = f"{PYTHON} {SPLATATLAS}/scripts/main_render.py --method {method} --source_path {src} --model_path {m_dir} --iteration 30000 --resolution {res}" + subprocess.run(render_cmd, shell=True) + + # 2. 评估 + renders = f"{m_dir}/renders_test_30000" + gt = f"{m_dir}/gt_test_30000" + depths = f"{m_dir}/depths_test_30000" + + eval_cmd = ( + f"{PYTHON} {SPLATATLAS}/ufd_evalkit/run_eval.py " + f"--method {method} --scene {scene} --render_dir {renders} " + f"--gt_dir {gt} --ply_path {ply_f} --output_json {metrics_f} " + f"--colmap_dir {src} --depth_dir {depths}" + ) + subprocess.run(eval_cmd, shell=True) + diff --git a/scripts/fix_missing_evals_v2.py b/scripts/fix_missing_evals_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..cceea14e4a86d425569eaa0eba64a4c32e565d98 --- /dev/null +++ b/scripts/fix_missing_evals_v2.py @@ -0,0 +1,70 @@ +import os, subprocess + +SPLATATLAS = "/root/autodl-tmp/SplatAtlas" +BASE = "/root/autodl-tmp/SplatAtlas/outputs" + +# 映射表：(环境路径, 额外 PYTHONPATH) +ENV_MAP = { + "vanilla_3dgs": ("base", "/root/autodl-tmp/3dgs_official"), + "analyticsplatting": ("/root/autodl-tmp/envs/ana_splatting", "/root/autodl-tmp/analyticsplatting_official"), + "erankgs": ("/root/autodl-tmp/envs/erank_gs", "/root/autodl-tmp/erank_gs"), + "pgsr": ("/root/autodl-tmp/envs/pgsr_v2", "/root/autodl-tmp/pgsr_official") +} + +def get_py_and_env(method): + if method not in ENV_MAP: + return "python", os.environ.copy() + + env_path, extra_py = ENV_MAP[method] + python_bin = "python" if env_path == "base" else f"{env_path}/bin/python" + + env = os.environ.copy() + env["PYTHONPATH"] = f"{SPLATATLAS}:{extra_py}:{env.get('PYTHONPATH', '')}" + # 缓解 LPIPS 评测时的 PyTorch 显存碎片化 OOM + env["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" + return python_bin, env + +targets = [ + ("vanilla_3dgs", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2), + ("analyticsplatting", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2), + ("analyticsplatting", "Lego", "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/Lego", 1), + ("erankgs", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2), + ("pgsr", "bonsai", "/root/autodl-tmp/dataset/360/bonsai", 2) +] +seeds = ["", "_seed1", "_seed2"] + +for method, scene, src, res in targets: + python_bin, custom_env = get_py_and_env(method) + + for suffix in seeds: + m_dir = f"{BASE}/{method}{suffix}_{scene}" + metrics_f = f"{m_dir}/metrics_test_iter30000.json" + ply_f = f"{m_dir}/point_cloud/iteration_30000/point_cloud.ply" + + if os.path.exists(metrics_f): + continue + + if not os.path.exists(ply_f): + print(f"[-] SKIPPING {method}{suffix}_{scene}: PLY not found") + continue + + print(f"\n[*] FIXING: {method}{suffix}_{scene}") + print(f" -> Using Python: {python_bin}") + + # 1. 渲染 (强制覆盖) + render_cmd = f"{python_bin} {SPLATATLAS}/scripts/main_render.py --method {method} --source_path {src} --model_path {m_dir} --iteration 30000 --resolution {res}" + subprocess.run(render_cmd, shell=True, env=custom_env) + + # 2. 评估 + renders = f"{m_dir}/renders_test_30000" + gt = f"{m_dir}/gt_test_30000" + depths = f"{m_dir}/depths_test_30000" + + eval_cmd = ( + f"{python_bin} {SPLATATLAS}/ufd_evalkit/run_eval.py " + f"--method {method} --scene {scene} --render_dir {renders} " + f"--gt_dir {gt} --ply_path {ply_f} --output_json {metrics_f} " + f"--colmap_dir {src} --depth_dir {depths}" + ) + subprocess.run(eval_cmd, shell=True, env=custom_env) + diff --git a/scripts/force_fix_drjohnson.py b/scripts/force_fix_drjohnson.py new file mode 100644 index 0000000000000000000000000000000000000000..ad1fd63f6f903bef7f456fd9d33491a62a2cdd9c --- /dev/null +++ b/scripts/force_fix_drjohnson.py @@ -0,0 +1,29 @@ +import os +from PIL import Image +from tqdm import tqdm + +def align_to_16(path): + print(f"🛠️ 正在对 {path} 进行物理对齐（居中裁剪至16倍数）...") + if not os.path.exists(path): + print("❌ 目录不存在！") + return + + files = sorted([f for f in os.listdir(path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) + for f in tqdm(files): + img_path = os.path.join(path, f) + with Image.open(img_path) as img: + w, h = img.size + # 计算最大的 16 倍数 + new_w = (w // 16) * 16 + new_h = (h // 16) * 16 + + if w != new_w or h != new_h: + # 居中裁剪逻辑 + left = (w - new_w) // 2 + top = (h - new_h) // 2 + right = left + new_w + bottom = top + new_h + img.crop((left, top, right, bottom)).save(img_path) + +# 执行物理修复 +align_to_16("/root/autodl-tmp/dataset/deepblending_clean/DrJohnson/images") diff --git a/scripts/functional_clustering.py b/scripts/functional_clustering.py new file mode 100644 index 0000000000000000000000000000000000000000..ba0e68c502c24845f838b858374e1912346f133e --- /dev/null +++ b/scripts/functional_clustering.py @@ -0,0 +1,66 @@ +import os +import glob +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import seaborn as sns +from PIL import Image +from umap import UMAP + +sns.set_theme(style="whitegrid") + +def load_ood_renders(method_dir): + ood_dir = os.path.join(method_dir, "renders_ood") + if not os.path.exists(ood_dir): + return None + images = sorted(glob.glob(os.path.join(ood_dir, "*.png"))) + if not images: + return None + + feature_vector = [] + for img_path in images: + img = Image.open(img_path).convert("RGB").resize((64, 64)) + feature_vector.append(np.array(img).flatten()) + + return np.concatenate(feature_vector) + +def main(): + outputs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "outputs")) + plots_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "plots")) + os.makedirs(plots_dir, exist_ok=True) + + features = [] + labels = [] + + for folder in os.listdir(outputs_dir): + method_dir = os.path.join(outputs_dir, folder) + if not os.path.isdir(method_dir) or folder == "pathology_archive": + continue + + vec = load_ood_renders(method_dir) + if vec is not None: + features.append(vec) + labels.append(folder.split('_')[0]) + + if not features: + return + + X = np.stack(features) + reducer = UMAP(n_components=2, random_state=42) + embedding = reducer.fit_transform(X) + + df = pd.DataFrame({ + 'UMAP_1': embedding[:, 0], + 'UMAP_2': embedding[:, 1], + 'Method': labels + }) + + plt.figure(figsize=(10, 8)) + sns.scatterplot(data=df, x='UMAP_1', y='UMAP_2', hue='Method', s=200, palette='tab10', edgecolor='k') + plt.title("Functional Space Convergence (OOD Renders)") + plt.tight_layout() + plt.savefig(os.path.join(plots_dir, "fig3_functional_manifold.png")) + plt.close() + +if __name__ == "__main__": + main() diff --git a/scripts/generate_slurm.py b/scripts/generate_slurm.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/scripts/main_render.py b/scripts/main_render.py new file mode 100644 index 0000000000000000000000000000000000000000..00c7575facd95f918a6056b42fadd7273e84fe17 --- /dev/null +++ b/scripts/main_render.py @@ -0,0 +1,91 @@ +import argparse +import os +import torch +import torchvision +import numpy as np +from tqdm import tqdm +import sys +import concurrent.futures + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from core.registry import METHOD_REGISTRY +import methods +from ufd_evalkit.geometric import depth_to_normal + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--method", type=str, required=True) + parser.add_argument("--source_path", type=str, required=True) + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument("--iteration", type=int, default=30000) + parser.add_argument("--resolution", type=int, default=-1) + return parser.parse_args() + +def save_np(data, path): np.save(path, data) + +def main(): + args = parse_args() + + # 动态从注册表获取模型，拒绝写死的 import + model_class = methods.load_method(args.method) + + dataset_config = {"source_path": args.source_path, "model_path": args.model_path, "resolution": args.resolution} + model = model_class(dataset_config, hyperparams={}) + + print(f"📥 [SplatAtlas] 正在请求变体自行加载资产... (Iteration: {args.iteration})") + try: + # 严格遵守契约：由 Wrapper 决定如何加载 .ply 或 MLP 权重 + model.load(args.model_path, args.iteration) + except Exception as e: + print(f"❌ 资产加载失败，请检查 [{args.method}] 的 load() 方法实现: {e}") + return + + print(f"📷 [SplatAtlas] 启动离线资产渲染 (异步 I/O 开启)") + + executor = concurrent.futures.ThreadPoolExecutor(max_workers=16) + futures = [] + + test_render_dir = os.path.join(args.model_path, f"renders_test_{args.iteration}") + test_depth_dir = os.path.join(args.model_path, f"depths_test_{args.iteration}") + test_normal_dir = os.path.join(args.model_path, f"normals_test_{args.iteration}") + test_gt_dir = os.path.join(args.model_path, f"gt_test_{args.iteration}") + for d in [test_render_dir, test_depth_dir, test_normal_dir, test_gt_dir]: os.makedirs(d, exist_ok=True) + + train_render_dir = os.path.join(args.model_path, f"renders_train_{args.iteration}") + for d in [train_render_dir]: os.makedirs(d, exist_ok=True) + + with torch.no_grad(): + test_cams = model.scene.getTestCameras() + if test_cams and isinstance(test_cams[0], list): + flat_cams = [c for sub in test_cams for c in sub] + else: + flat_cams = test_cams + for idx, view in enumerate(tqdm(flat_cams, desc="Rendering Test")): + render_pkg = model.render(view) + rendering = render_pkg.get("render", render_pkg.get("image")).detach().cpu() + depth = render_pkg.get("depth", None) + + futures.append(executor.submit(torchvision.utils.save_image, rendering, os.path.join(test_render_dir, f'{idx:05d}.png'))) + futures.append(executor.submit(torchvision.utils.save_image, view.original_image.cpu(), os.path.join(test_gt_dir, f'{idx:05d}.png'))) + + if depth is not None: + depth_cpu = depth.detach().cpu() + normal = render_pkg.get("normal", depth_to_normal(depth)).detach().cpu() + futures.append(executor.submit(save_np, depth_cpu.numpy(), os.path.join(test_depth_dir, f'{idx:05d}.npy'))) + futures.append(executor.submit(torchvision.utils.save_image, normal, os.path.join(test_normal_dir, f'{idx:05d}.png'))) + + for idx, view in enumerate(tqdm(model.scene.getTrainCameras(), desc="Rendering Train")): + render_pkg = model.render(view) + rendering = render_pkg.get("render", render_pkg.get("image")).detach().cpu() + futures.append(executor.submit(torchvision.utils.save_image, rendering, os.path.join(train_render_dir, f'{idx:05d}.png'))) + + print("⏳ [Async I/O] GPU 渲染完毕！正在等待 CPU 将所有数据落盘...") + concurrent.futures.wait(futures) + executor.shutdown() + + flag_path = os.path.join(args.model_path, f"render_complete_{args.iteration}.flag") + with open(flag_path, "w") as f: f.write("Completed by Async Engine") + print("✅ 完美落盘，标记成功！") + +if __name__ == "__main__": + main() diff --git a/scripts/main_train.py b/scripts/main_train.py new file mode 100644 index 0000000000000000000000000000000000000000..c2ae52881951de18c452e4c4c84e3eab8281f826 --- /dev/null +++ b/scripts/main_train.py @@ -0,0 +1,231 @@ + +import argparse +import os +import json +import torch +import torchvision +import numpy as np +from tqdm import tqdm +from torch.utils.tensorboard import SummaryWriter +import sys +import time +import concurrent.futures + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from ufd_evalkit.geometric import depth_to_normal +import methods +from core.registry import METHOD_REGISTRY + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--method", type=str, required=True) + parser.add_argument("--source_path", type=str, required=True) + parser.add_argument("--model_path", type=str, required=True) + parser.add_argument("--iterations", type=int, default=30000) + parser.add_argument("--save_iterations", nargs="+", type=int, default=[5000, 10000, 20000, 30000]) + parser.add_argument("--resolution", type=int, default=-1) + parser.add_argument("--track_decoupling", action="store_true") + parser.add_argument("--cap_gaussians", type=int, default=None) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--skip_render", action="store_true") + return parser.parse_args() + +def save_np(data, path): np.save(path, data) + +def in_memory_render(model, model_path, iteration): + print("\n[Memory Fusion] Flushing VRAM fragments...") + torch.cuda.empty_cache() + + executor = concurrent.futures.ThreadPoolExecutor(max_workers=16) + futures = [] + + with torch.no_grad(): + test_render_dir = os.path.join(model_path, f"renders_test_{iteration}") + test_depth_dir = os.path.join(model_path, f"depths_test_{iteration}") + test_normal_dir = os.path.join(model_path, f"normals_test_{iteration}") + test_gt_dir = os.path.join(model_path, f"gt_test_{iteration}") + for d in [test_render_dir, test_depth_dir, test_normal_dir, test_gt_dir]: os.makedirs(d, exist_ok=True) + + print(f"Launching async full-dimensional rendering | Test Set") + for idx, view in enumerate(tqdm(model.scene.getTestCameras(), desc="Rendering Test")): + render_pkg = model.render(view) + rendering = render_pkg.get("render", render_pkg.get("image")).detach().cpu() + depth = render_pkg.get("depth", None) + + futures.append(executor.submit(torchvision.utils.save_image, rendering, os.path.join(test_render_dir, f'{idx:05d}.png'))) + futures.append(executor.submit(torchvision.utils.save_image, view.original_image.cpu(), os.path.join(test_gt_dir, f'{idx:05d}.png'))) + + if depth is not None: + depth_cpu = depth.detach().cpu() + normal = render_pkg.get("normal", depth_to_normal(view, depth)).detach().cpu() + futures.append(executor.submit(save_np, depth_cpu.numpy(), os.path.join(test_depth_dir, f'{idx:05d}.npy'))) + futures.append(executor.submit(torchvision.utils.save_image, normal, os.path.join(test_normal_dir, f'{idx:05d}.png'))) + + train_render_dir = os.path.join(model_path, f"renders_train_{iteration}") + train_depth_dir = os.path.join(model_path, f"depths_train_{iteration}") + for d in [train_render_dir, train_depth_dir]: os.makedirs(d, exist_ok=True) + + print(f"Launching async full-dimensional rendering | Train Set") + for idx, view in enumerate(tqdm(model.scene.getTrainCameras(), desc="Rendering Train")): + render_pkg = model.render(view) + rendering = render_pkg.get("render", render_pkg.get("image")).detach().cpu() + futures.append(executor.submit(torchvision.utils.save_image, rendering, os.path.join(train_render_dir, f'{idx:05d}.png'))) + + depth = render_pkg.get("depth", None) + if depth is not None: + futures.append(executor.submit(save_np, depth.detach().cpu().numpy(), os.path.join(train_depth_dir, f'{idx:05d}.npy'))) + + print("[Async I/O] GPU rendering complete. Waiting for CPU flush...") + concurrent.futures.wait(futures) + + error_log_path = os.path.join(model_path, "render_error.log") + has_error = False + for future in futures: + try: + future.result() + except Exception as e: + if not has_error: + print(f"\n[CRITICAL ERROR] Async render thread crashed. See {error_log_path}") + has_error = True + with open(error_log_path, "a") as err_f: + err_f.write(f"Exception: {str(e)}\n") + + executor.shutdown() + + if not has_error: + flag_path = os.path.join(model_path, f"render_complete_{iteration}.flag") + with open(flag_path, "w") as f: f.write("Completed by Async Engine") + print("Flush complete. Flag written.") + else: + print("[WARNING] Render flag not written due to errors.") + +def main(): + import random + args = parse_args() + random.seed(args.seed); np.random.seed(args.seed) + torch.manual_seed(args.seed); torch.cuda.manual_seed_all(args.seed) + + os.makedirs(args.model_path, exist_ok=True) + tb_writer = SummaryWriter(args.model_path) + + dataset_config = {"source_path": args.source_path, "model_path": args.model_path, "resolution": args.resolution} + model = methods.load_method(args.method)(dataset_config, {"track_decoupling": args.track_decoupling, "cap_gaussians": args.cap_gaussians}) + + print(f"[SplatAtlas] Starting training | Method: {args.method} | Resolution: {args.resolution}") + + num_train = len(model.scene.getTrainCameras()) + num_test = len(model.scene.getTestCameras()) if hasattr(model.scene, "getTestCameras") else 0 + split_info = ( + f"========================================\n" + f"[Dataset Topology Checker]\n" + f"Method: {args.method}\n" + f"Train Cameras: {num_train}\n" + f"Test Cameras: {num_test}\n" + f"Leakage Status: {'SAFE' if num_test > 0 else 'DANGER (NO TEST SET DETECTED)'}\n" + f"========================================" + ) + print(split_info) + + with open(os.path.join(args.model_path, "dataset_split.log"), "w") as f: + f.write(split_info + "\n") + + cam_json_path = os.path.join(args.model_path, "cameras.json") + with open(cam_json_path, "w") as f: + json.dump([{"id": i, "center": c.camera_center.detach().cpu().numpy().tolist(), "view_dir": c.world_view_transform[:3, 2].detach().cpu().numpy().tolist()} for i, c in enumerate([c for item in (model.scene.getTrainCameras().values() if isinstance(model.scene.getTrainCameras(), dict) else model.scene.getTrainCameras()) for c in (item if isinstance(item, list) else [item])])], f, indent=4) + + t_train_start = time.time() + timing_path = os.path.join(args.model_path, "timing.json") + with open(timing_path, "w") as f: + json.dump({"train_start_unix": t_train_start, "status": "running"}, f) + + cumulative_gpu_time = 0.0 + for step in range(1, args.iterations + 1): + t_step_start = time.time() + + stats_or_tuple = model.train_iteration(step) + + if isinstance(stats_or_tuple, tuple): + stats, histograms = stats_or_tuple + else: + stats, histograms = stats_or_tuple, {} + + step_elapsed = time.time() - t_step_start + cumulative_gpu_time += step_elapsed + stats['cumulative_gpu_time_sec'] = cumulative_gpu_time + stats['step_time_ms'] = step_elapsed * 1000.0 + if step % 100 == 0: + print(f"[TIMING] step={step} step_time_ms={stats.get('step_time_ms',0):.2f} " + f"iter_time_ms={stats.get('iter_time_ms',0):.2f} " + f"peak_vram_GB={stats.get('peak_vram_GB',0):.3f}", flush=True) + + if step % 100 == 0: + print(f"\n{'='*40}") + print(f"Step {step:05d} Runtime Monitor") + print(f"{'='*40}") + + print("[Scalars]") + for key, val in stats.items(): + if val is not None: + tb_writer.add_scalar(f"train/{key}", val, step) + if isinstance(val, float): + print(f" |- {key}: {val:.6f}") + else: + print(f" |- {key}: {val}") + + if histograms: + print("\n[Histograms]") + for dist_name, tensor_data in histograms.items(): + if not torch.isnan(tensor_data).any(): + try: + arr = tensor_data.float().cpu().numpy().flatten() + arr = arr[np.isfinite(arr)] + if len(arr) > 0: + counts, bin_edges = np.histogram(arr.astype(np.float64), bins=100) + tb_writer.add_histogram_raw( + f'Distributions/{dist_name}', + min=float(arr.min()), max=float(arr.max()), + num=len(arr), + sum=float(arr.sum()), + sum_squares=float((arr.astype(np.float64)**2).sum()), + bucket_limits=bin_edges[1:].tolist(), + bucket_counts=counts.tolist(), + global_step=step + ) + except Exception as e: + print(f" |- histogram write failed: {e}") + t_min = tensor_data.min().item() + t_max = tensor_data.max().item() + t_mean = tensor_data.mean().item() + shape_str = "x".join(map(str, tensor_data.shape)) + print(f" |- {dist_name}:") + print(f" | - Shape: [{shape_str}], Mean: {t_mean:.4f}, Min: {t_min:.4f}, Max: {t_max:.4f}") + else: + print(f" |- {dist_name}: contains NaN, skipped.") + + if step in args.save_iterations: + model.save(args.model_path, step) + print(f"\n[CheckPoint] Assets saved to: {args.model_path} (Step {step})") + + t_train_end = time.time() + total_training_seconds = t_train_end - t_train_start + + with open(timing_path, "w") as f: + json.dump({ + "train_start_unix": t_train_start, + "train_end_unix": t_train_end, + "total_training_seconds": total_training_seconds, + "total_training_minutes": total_training_seconds / 60.0, + "cumulative_gpu_time_seconds": cumulative_gpu_time, + "iterations": args.iterations, + "status": "complete" + }, f, indent=4) + + print(f"\n[Timing] Training complete. Total: {total_training_seconds:.1f}s / {total_training_seconds/60.0:.2f}min") + print(f"[Timing] timing.json written to {timing_path}") + + tb_writer.close() + if not args.skip_render: + in_memory_render(model, args.model_path, args.iterations) + +if __name__ == "__main__": + main() diff --git a/scripts/patch_json.py b/scripts/patch_json.py new file mode 100644 index 0000000000000000000000000000000000000000..b869b34aef7e5bcfb4bebcb4082c948068e36f27 --- /dev/null +++ b/scripts/patch_json.py @@ -0,0 +1,34 @@ +import os + +def patch_json_dump(target_file): + if not os.path.exists(target_file): + return + with open(target_file, "r") as f: + code = f.read() + + if "class NpEncoder(json.JSONEncoder):" in code: + print(f"Already patched: {target_file}") + return + + # 注入万能类型转换器 + encoder_code = """import json + +class NpEncoder(json.JSONEncoder): + def default(self, obj): + if hasattr(obj, 'item'): + return obj.item() + return super().default(obj) +""" + code = code.replace("import json", encoder_code) + + # 替换所有的 dump 调用 + code = code.replace("json.dump(stats, f, indent=4)", "json.dump(stats, f, indent=4, cls=NpEncoder)") + code = code.replace("json.dump(metrics, f, indent=4)", "json.dump(metrics, f, indent=4, cls=NpEncoder)") + code = code.replace("json.dump(results, f, indent=4)", "json.dump(results, f, indent=4, cls=NpEncoder)") + + with open(target_file, "w") as f: + f.write(code) + print(f"Patched JSON serialization in {target_file}") + +patch_json_dump("/root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py") +patch_json_dump("/root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py") diff --git a/scripts/patch_oom.py b/scripts/patch_oom.py new file mode 100644 index 0000000000000000000000000000000000000000..569c027f69c34821a4674e6f177be1e71ee92f9e --- /dev/null +++ b/scripts/patch_oom.py @@ -0,0 +1,26 @@ +import os + +target_file = "/root/autodl-tmp/SplatAtlas/methods/wrapper_2dgs.py" +with open(target_file, "r") as f: + code = f.read() + +old_block = """ chunk_size = 2000 + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + pts_chunk = query_points[i:end] + dist_sq = torch.cdist(pts_chunk, xyz, p=2).pow(2)""" + +new_block = """ N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // N_gaussians) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + pts_chunk = query_points[i:end] + dist_sq = torch.cdist(pts_chunk, xyz, p=2).pow(2)""" + +if old_block in code: + code = code.replace(old_block, new_block) + with open(target_file, "w") as f: + f.write(code) + print("Patch applied successfully.") +else: + print("Block not found or already patched.") diff --git a/scripts/phase1_validation/dump_tnt_worst.py b/scripts/phase1_validation/dump_tnt_worst.py new file mode 100644 index 0000000000000000000000000000000000000000..2d20492fe6606f02fb365c863a0f0f6d52eb00b4 --- /dev/null +++ b/scripts/phase1_validation/dump_tnt_worst.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python +"""Dump <20dB render_vs_render cam pairs for manual inspection.""" +import os, sys, inspect, math +import numpy as np +import torch +from PIL import Image, ImageOps +from plyfile import PlyData + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') +import gsplat + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +DUMP_ROOT = "/tmp/tnt_diagnose" +SCENE = "lighthouse" +THRESHOLD = 20.0 # render_vs_render PSNR < THRESHOLD 的 cam 会被 dump + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + a = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: a.append(source_path) + elif i == 1: a.append(img_dir) + elif k == "eval": a.append(True) + elif k == "train_test_exp": a.append(False) + else: a.append(p.default if p.default != inspect.Parameter.empty else "") + return a + + +def to_u8(arr): + return np.clip(arr * 255, 0, 255).astype(np.uint8) if arr.dtype != np.uint8 else arr + + +def hstack_panels(arrs, gap=20): + imgs = [to_u8(a) for a in arrs] + H = imgs[0].shape[0] + W = sum(i.shape[1] for i in imgs) + gap * (len(imgs) - 1) + out = np.ones((H, W, 3), dtype=np.uint8) * 240 + x = 0 + for img in imgs: + out[:, x:x + img.shape[1]] = img + x += img.shape[1] + gap + return out + + +def main(): + source_path = os.path.join(DATASET_ROOT, SCENE) + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{SCENE}") + dump = os.path.join(DUMP_ROOT, SCENE) + os.makedirs(dump, exist_ok=True) + # 清空旧内容 + for f in os.listdir(dump): + os.remove(os.path.join(dump, f)) + + img_dir = "images_2"; resolution = 2 + + from scene.dataset_readers import readColmapSceneInfo + from utils.graphics_utils import getWorld2View2 + from utils.general_utils import PILtoTorch + + scene_info = readColmapSceneInfo(*build_args(source_path, img_dir)) + test_cams = scene_info.test_cameras + print(f"test cams: {len(test_cams)}") + + # reverse-lookup: real image_name → native_idx + native_render_dir = os.path.join(cell, "renders_test_30000") + if os.path.isdir(os.path.join(native_render_dir, "renders")): + native_render_dir = os.path.join(native_render_dir, "renders") + native_gt_dir = os.path.join(cell, "gt_test_30000") + native_renders = sorted([f for f in os.listdir(native_render_dir) + if f.lower().endswith(('.png', '.jpg'))]) + native_gts = sorted([f for f in os.listdir(native_gt_dir) + if f.lower().endswith(('.png', '.jpg'))]) + + THUMB = 128 + i2_files = sorted([f for f in os.listdir(os.path.join(source_path, img_dir)) + if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) + i2_thumbs = np.stack([ + np.asarray(Image.open(os.path.join(source_path, img_dir, f)).convert('RGB') + .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0 + for f in i2_files]) + real_to_native_idx = {} + for nidx, ng in enumerate(native_gts): + t = np.asarray(Image.open(os.path.join(native_gt_dir, ng)).convert('RGB') + .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0 + mse = ((i2_thumbs - t) ** 2).mean(axis=(1, 2, 3)) + real_to_native_idx[i2_files[int(np.argmin(mse))]] = nidx + + # PLY + pd = PlyData.read(os.path.join(cell, "point_cloud", "iteration_30000", "point_cloud.ply")) + vv = pd['vertex'] + device = torch.device("cuda") + def t32(x): return torch.tensor(x, dtype=torch.float32, device=device) + means = t32(np.stack((vv['x'], vv['y'], vv['z']), -1)) + quats = t32(np.stack((vv['rot_0'], vv['rot_1'], vv['rot_2'], vv['rot_3']), -1)) + scales = torch.exp(t32(np.stack((vv['scale_0'], vv['scale_1'], vv['scale_2']), -1))) + opacities = torch.sigmoid(t32(np.asarray(vv['opacity']))) + f_dc = t32(np.stack((vv['f_dc_0'], vv['f_dc_1'], vv['f_dc_2']), -1)).unsqueeze(1) + f_rest = t32(np.stack([vv[f'f_rest_{i}'] for i in range(45)], -1)).view(-1, 3, 15).transpose(1, 2) + shs = torch.cat([f_dc, f_rest], dim=1) + bg_color = torch.tensor([0., 0., 0.], device=device) + + summary = [] + for i, c in enumerate(test_cams): + w = int(round(c.width / resolution)); h = int(round(c.height / resolution)) + viewmat = t32(getWorld2View2(np.array(c.R), np.array(c.T))).unsqueeze(0) + fx = w / (2 * math.tan(c.FovX / 2)); fy = h / (2 * math.tan(c.FovY / 2)) + K = t32(np.array([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]])).unsqueeze(0) + + with torch.no_grad(): + colors, _, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, + colors=shs, viewmats=viewmat, Ks=K, width=w, height=h, + sh_degree=3, packed=True, render_mode='RGB', + backgrounds=bg_color.unsqueeze(0)) + render = colors[0].clamp(0, 1).cpu().numpy() # (H, W, 3) float + + if c.image_name not in real_to_native_idx: + continue + nidx = real_to_native_idx[c.image_name] + nr = np.asarray(Image.open(os.path.join(native_render_dir, native_renders[nidx])) + .convert('RGB'), dtype=np.float32) / 255.0 + ng = np.asarray(Image.open(os.path.join(native_gt_dir, native_gts[nidx])) + .convert('RGB'), dtype=np.float32) / 255.0 + # 统一尺寸到 render + if nr.shape != render.shape: + th, tw = render.shape[:2] + nr = np.asarray(Image.fromarray(to_u8(nr)).resize((tw, th), Image.LANCZOS), dtype=np.float32) / 255.0 + ng = np.asarray(Image.fromarray(to_u8(ng)).resize((tw, th), Image.LANCZOS), dtype=np.float32) / 255.0 + + vs_mse = ((render - nr) ** 2).mean() + vs_psnr = 10 * math.log10(1.0 / max(vs_mse, 1e-10)) + ours_p = 10 * math.log10(1.0 / max(((render - ng) ** 2).mean(), 1e-10)) + nat_p = 10 * math.log10(1.0 / max(((nr - ng) ** 2).mean(), 1e-10)) + + mark = "DUMP" if vs_psnr < THRESHOLD else "skip" + print(f" {i:>3} {c.image_name:>12} vs={vs_psnr:6.2f} ours={ours_p:6.2f} nat={nat_p:6.2f} {mark}") + + if vs_psnr < THRESHOLD: + base = os.path.splitext(c.image_name)[0] + Image.fromarray(to_u8(ng)).save(os.path.join(dump, f"{base}_1_gt.png")) + Image.fromarray(to_u8(render)).save(os.path.join(dump, f"{base}_2_ours.png")) + Image.fromarray(to_u8(nr)).save(os.path.join(dump, f"{base}_3_native.png")) + diff5 = np.clip(np.abs(render - nr) * 5, 0, 1) + diff20 = np.clip(np.abs(render - nr) * 20, 0, 1) + Image.fromarray(to_u8(diff5)).save(os.path.join(dump, f"{base}_4_diff5x.png")) + Image.fromarray(to_u8(diff20)).save(os.path.join(dump, f"{base}_5_diff20x.png")) + # 4-panel: GT | OURS | NATIVE | DIFF*5 + panel = hstack_panels([ng, render, nr, diff5]) + Image.fromarray(panel).save(os.path.join(dump, f"{base}_panel.png")) + summary.append((i, c.image_name, vs_psnr, ours_p, nat_p)) + + # summary + sp = os.path.join(dump, "summary.txt") + with open(sp, "w") as f: + f.write(f"Scene: {SCENE}\n") + f.write(f"Threshold: render_vs_render < {THRESHOLD} dB\n") + f.write(f"Dumped : {len(summary)} cams\n\n") + f.write(f"{'idx':>3} {'image_name':>14} {'vs_native':>10} {'ours_psnr':>10} {'nat_psnr':>9}\n") + f.write("-" * 60 + "\n") + for s in sorted(summary, key=lambda x: x[2]): + f.write(f"{s[0]:>3} {s[1]:>14} {s[2]:>10.2f} {s[3]:>10.2f} {s[4]:>9.2f}\n") + f.write("\nPer-cam files:\n") + f.write(" _1_gt.png GT (Native 存盘)\n") + f.write(" _2_ours.png gsplat render\n") + f.write(" _3_native.png diff-gaussian-rasterization render\n") + f.write(" _4_diff5x.png |ours-native| * 5 (彩色)\n") + f.write(" _5_diff20x.png |ours-native| * 20 (极端放大)\n") + f.write(" _panel.png GT | OURS | NATIVE | DIFF*5 横排\n") + + print(f"\nDumped {len(summary)} cams to:") + print(f" {dump}/") + print(f"Summary: {sp}") + print(f"\n拉回本地查看：") + print(f" scp -P -r root@:{dump} ./tnt_{SCENE}_diag") + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/probe_tnt.py b/scripts/phase1_validation/probe_tnt.py new file mode 100644 index 0000000000000000000000000000000000000000..a8392bf813c8ab40a5577eb4b4aa5dfda5891cd6 --- /dev/null +++ b/scripts/phase1_validation/probe_tnt.py @@ -0,0 +1,227 @@ +#!/usr/bin/env python +"""T&T root-cause probe — read-only, symmetric PASS vs FAIL.""" +import os, sys, inspect, struct +import numpy as np +from PIL import Image, ImageOps +from plyfile import PlyData + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")] + + +def sec(title): + print("\n" + "=" * 70) + print(f" {title}") + print("=" * 70) + + +def build_args(source_path, img_dir): + """与 render_single.py 完全一致的 args 构造逻辑""" + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + args_list = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: args_list.append(source_path) + elif i == 1: args_list.append(img_dir) + elif k == "eval": args_list.append(True) + elif k == "train_test_exp": args_list.append(False) + else: args_list.append(p.default if p.default != inspect.Parameter.empty else "") + return sig, args_list + + +# ---------- Probe 1 ---------- +def probe_signature(): + sec("PROBE 1 — readColmapSceneInfo signature & args mapping") + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + print(f"Full signature: {sig}") + print() + print("按 render_single.py 的位置策略喂参 (source_path='', img_dir=''):") + for i, (k, p) in enumerate(sig.parameters.items()): + default = p.default if p.default != inspect.Parameter.empty else "" + if i == 0: fed = "" + elif i == 1: fed = "" + elif k == "eval": fed = "True" + elif k == "train_test_exp": fed = "False" + else: fed = f"default({default})" + warn = "" + if k in ("eval", "train_test_exp") and i < 2: + warn = " <<< WARNING: positional slot collides with keyword logic" + print(f" [{i}] {k:<20} default={default!s:<20} -> fed: {fed}{warn}") + + +# ---------- Probe 2 ---------- +def probe_native_gt(scene): + sec(f"PROBE 2 — Native GT manifest [{scene}]") + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + candidates = [ + os.path.join(cell, "gt_test_30000"), + os.path.join(cell, "renders_test_30000", "gt"), + os.path.join(cell, "test", "ours_30000", "gt"), + ] + gt_dir = next((c for c in candidates if os.path.isdir(c)), None) + if gt_dir is None: + print(f"[!] GT dir not found. Tried: {candidates}") + return None + files = sorted([f for f in os.listdir(gt_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))]) + print(f"GT dir : {gt_dir}") + print(f"Count : {len(files)}") + print(f"First 5: {files[:5]}") + print(f"Last 3: {files[-3:]}") + if files: + img = Image.open(os.path.join(gt_dir, files[0])) + print(f"Sample dims (first GT): {img.size} (W x H)") + return set(os.path.splitext(f)[0] for f in files) + + +# ---------- Probe 3 ---------- +def probe_our_split(scene): + sec(f"PROBE 3 — Our eval=True split [{scene}]") + from scene.dataset_readers import readColmapSceneInfo + source_path = os.path.join(DATASET_ROOT, scene) + img_dir = "images_2" if os.path.exists(os.path.join(source_path, "images_2")) else "images" + print(f"Picked img_dir: {img_dir}") + sig, args_list = build_args(source_path, img_dir) + print(f"Args positionally passed:") + for i, (k, _) in enumerate(sig.parameters.items()): + print(f" [{i}] {k} = {args_list[i]!r}") + try: + scene_info = readColmapSceneInfo(*args_list) + except Exception as e: + print(f"[!] readColmapSceneInfo crashed: {type(e).__name__}: {e}") + return None + tr, te = scene_info.train_cameras, scene_info.test_cameras + print(f"Train cams : {len(tr)}") + print(f"Test cams : {len(te)}") + if te: + names = sorted(c.image_name for c in te) + print(f"First 5 test names: {names[:5]}") + print(f"Last 3 test names: {names[-3:]}") + c0 = te[0] + print(f"First test cam dims: W={c0.width} H={c0.height} FovX={c0.FovX:.4f} FovY={c0.FovY:.4f}") + return set(os.path.splitext(c.image_name)[0] for c in te) + + +# ---------- Probe 4 ---------- +def probe_split_diff(scene, native_set, ours_set): + sec(f"PROBE 4 — Split diff [{scene}]") + if native_set is None or ours_set is None: + print("[!] skipped (missing input)") + return + only_n = sorted(native_set - ours_set) + only_o = sorted(ours_set - native_set) + print(f"|Native|={len(native_set)} |Ours|={len(ours_set)} |intersection|={len(native_set & ours_set)}") + print(f"In Native only ({len(only_n)}): {only_n[:5]}") + print(f"In Ours only ({len(only_o)}): {only_o[:5]}") + if not only_n and not only_o: + print("[OK] Test split fully aligned.") + else: + print("[!!] SPLIT MISMATCH — eval=True 没把我们带到和 Native 一样的测试集上") + + +# ---------- Probe 5 ---------- +def probe_camera_model(scene): + sec(f"PROBE 5 — COLMAP camera model [{scene}]") + src = os.path.join(DATASET_ROOT, scene, "sparse", "0") + txt, binf = os.path.join(src, "cameras.txt"), os.path.join(src, "cameras.bin") + if os.path.exists(txt): + with open(txt) as f: + for line in f: + if line.startswith('#') or not line.strip(): continue + p = line.strip().split() + print(f" id={p[0]} model={p[1]} size={p[2]}x{p[3]} params={p[4:]}") + break + elif os.path.exists(binf): + MODELS = {0:"SIMPLE_PINHOLE",1:"PINHOLE",2:"SIMPLE_RADIAL",3:"RADIAL", + 4:"OPENCV",5:"OPENCV_FISHEYE",6:"FULL_OPENCV",7:"FOV", + 8:"SIMPLE_RADIAL_FISHEYE",9:"RADIAL_FISHEYE",10:"THIN_PRISM_FISHEYE"} + with open(binf, "rb") as f: + num = struct.unpack("50 : {(z>50).sum():>8} ({100*(z>50).mean():.2f}%)") + print(f"Gaussians Z>100 : {(z>100).sum():>8} ({100*(z>100).mean():.2f}%)") + print(f"Gaussians Z>500 : {(z>500).sum():>8} ({100*(z>500).mean():.2f}%)") + + +def main(): + probe_signature() + for scene, label in SCENES: + print(f"\n\n{'#'*70}\n# SCENE: {scene} [{label}]\n{'#'*70}") + native = probe_native_gt(scene) + ours = probe_our_split(scene) + probe_split_diff(scene, native, ours) + probe_camera_model(scene) + probe_exif(scene) + probe_ply_z(scene) + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/probe_tnt2.py b/scripts/phase1_validation/probe_tnt2.py new file mode 100644 index 0000000000000000000000000000000000000000..84db1729ea504e3ab5ff75aace50fc866b410849 --- /dev/null +++ b/scripts/phase1_validation/probe_tnt2.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python +"""T&T round-2 probe — pinpoint remaining suspects at pixel level.""" +import os, sys, inspect, struct +import numpy as np +from PIL import Image + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")] + + +def sec(title): + print("\n" + "=" * 70) + print(f" {title}") + print("=" * 70) + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + args = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: args.append(source_path) + elif i == 1: args.append(img_dir) + elif k == "eval": args.append(True) + elif k == "train_test_exp": args.append(False) + else: args.append(p.default if p.default != inspect.Parameter.empty else "") + return args + + +def psnr_pil(a_pil, b_pil): + """统一 resize 到较小者，RGB 对齐 PSNR。""" + if a_pil.size != b_pil.size: + tgt = (min(a_pil.size[0], b_pil.size[0]), + min(a_pil.size[1], b_pil.size[1])) + a_pil = a_pil.resize(tgt, Image.LANCZOS) + b_pil = b_pil.resize(tgt, Image.LANCZOS) + a = np.asarray(a_pil.convert("RGB"), dtype=np.float32) / 255.0 + b = np.asarray(b_pil.convert("RGB"), dtype=np.float32) / 255.0 + mse = float(((a - b) ** 2).mean()) + return (100.0 if mse == 0 else 10 * np.log10(1.0 / mse)), mse + + +# ---------- Probe 8: PINHOLE full params (cx, cy) ---------- +def probe_camera_params(scene): + sec(f"PROBE 8 — COLMAP PINHOLE full params [{scene}]") + binf = os.path.join(DATASET_ROOT, scene, "sparse", "0", "cameras.bin") + if not os.path.exists(binf): + print(f"[!] no {binf}") + return + MODELS = {0:("SIMPLE_PINHOLE",3), 1:("PINHOLE",4), 2:("SIMPLE_RADIAL",4), + 3:("RADIAL",5), 4:("OPENCV",8), 5:("OPENCV_FISHEYE",8), + 6:("FULL_OPENCV",12), 7:("FOV",5), + 8:("SIMPLE_RADIAL_FISHEYE",4), 9:("RADIAL_FISHEYE",5), + 10:("THIN_PRISM_FISHEYE",12)} + with open(binf, "rb") as f: + num = struct.unpack(" 0.5 or off_y > 0.5: + print(f" [!!] PRINCIPAL POINT OFFSET >0.5% — our K矩阵用 w/2,h/2 忽略了这个偏移！") + else: + print(f" [OK] 主点在中心附近 (<0.5%偏移)") + + +# ---------- Probe 9: Native GT vs Ours GT pixel match ---------- +def probe_gt_match(scene): + sec(f"PROBE 9 — GT content alignment (Native-saved vs images_2/) [{scene}]") + from scene.dataset_readers import readColmapSceneInfo + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + gt_dir = os.path.join(cell, "gt_test_30000") + if not os.path.isdir(gt_dir): + gt_dir = os.path.join(cell, "renders_test_30000", "gt") + if not os.path.isdir(gt_dir): + print(f"[!] 没找到 Native GT 目录") + return None + native_gts = sorted([f for f in os.listdir(gt_dir) if f.lower().endswith(('.png','.jpg','.jpeg'))]) + + source_path = os.path.join(DATASET_ROOT, scene) + img_dir = "images_2" if os.path.exists(os.path.join(source_path, "images_2")) else "images" + scene_info = readColmapSceneInfo(*build_args(source_path, img_dir)) + our_test = scene_info.test_cameras + + print(f"Native GT dir : {gt_dir} count={len(native_gts)}") + print(f"Our test cams : count={len(our_test)} (first={our_test[0].image_name})") + if len(native_gts) != len(our_test): + print(f"[!!] COUNT MISMATCH — 直接说明两组 test 根本不同") + return our_test, native_gts, gt_dir + + print(f"\n假说 A: Native 按 idx rename GT，内容 == Ours。验证方式：对齐 idx 比较像素。") + print(f"{'idx':>3} {'native_file':>15} {'ours_name':>15} {'nat_size':>11} {'our_size':>11} {'PSNR':>8} {'MSE':>10}") + print("-" * 78) + psnrs = [] + for i in range(len(native_gts)): + nat_pil = Image.open(os.path.join(gt_dir, native_gts[i])) + our_path = our_test[i].image_path + if not os.path.exists(our_path): + our_path = os.path.join(source_path, img_dir, our_test[i].image_name) + our_pil = Image.open(our_path) + p, m = psnr_pil(nat_pil, our_pil) + psnrs.append(p) + if i < 5: + print(f"{i:>3} {native_gts[i]:>15} {our_test[i].image_name:>15} " + f"{str(nat_pil.size):>11} {str(our_pil.size):>11} {p:>8.2f} {m:>10.6f}") + + arr = np.array(psnrs) + print(f"\nAll {len(psnrs)} pairs: mean={arr.mean():.2f} min={arr.min():.2f} max={arr.max():.2f}") + print(f" >40 dB: {(arr>40).sum()}/{len(arr)} 20-40 dB: {((arr>=20)&(arr<=40)).sum()}/{len(arr)} <20 dB: {(arr<20).sum()}/{len(arr)}") + + if arr.mean() > 50: + print("[A 验证通过] GT 内容完全对齐。FAIL 原因在 renderer/pose/intrinsics 层面。") + elif arr.mean() > 30: + print("[A 弱成立] GT 内容近似一致 (30-50 dB) — 可能有 downsample 滤波器差异，但不足以解释 -1.24 dB") + else: + print("[A 被推翻] Native 和 Ours 拿到的是不同的图 —— 就是 test split 真的不一样。") + return our_test, native_gts, gt_dir + + +# ---------- Probe 10: Native baseline 自洽性 sanity ---------- +def probe_native_sanity(scene, our_test, native_gts, gt_dir): + sec(f"PROBE 10 — Native render vs Native GT (reproduces baseline) [{scene}]") + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + candidates = [ + os.path.join(cell, "renders_test_30000", "renders"), + os.path.join(cell, "renders_test_30000"), + os.path.join(cell, "test", "ours_30000", "renders"), + ] + renders_dir = next((c for c in candidates if os.path.isdir(c) + and any(f.lower().endswith(('.png','.jpg')) for f in os.listdir(c))), None) + if renders_dir is None: + print("[!] 没找到 Native renders 目录，跳过") + return + native_renders = sorted([f for f in os.listdir(renders_dir) if f.lower().endswith(('.png','.jpg'))]) + print(f"Native renders: {renders_dir} count={len(native_renders)}") + if not native_renders: + return + + psnrs = [] + for i in range(min(len(native_renders), len(native_gts))): + r = Image.open(os.path.join(renders_dir, native_renders[i])) + g = Image.open(os.path.join(gt_dir, native_gts[i])) + p, _ = psnr_pil(r, g) + psnrs.append(p) + arr = np.array(psnrs) + print(f"Native_render vs Native_GT mean PSNR = {arr.mean():.4f} dB " + f"(stored baseline should match)") + print(f" → 这个值就是 metrics_test_iter30000.json 里的 Native baseline。复核一下。") + + +def main(): + for scene, label in SCENES: + print(f"\n\n{'#'*70}\n# SCENE: {scene} [{label}]\n{'#'*70}") + probe_camera_params(scene) + result = probe_gt_match(scene) + if result: + our_test, native_gts, gt_dir = result + probe_native_sanity(scene, our_test, native_gts, gt_dir) + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/probe_tnt3.py b/scripts/phase1_validation/probe_tnt3.py new file mode 100644 index 0000000000000000000000000000000000000000..c700cf355d6c2e6a768cdbf081f0c6ca882c11ec --- /dev/null +++ b/scripts/phase1_validation/probe_tnt3.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python +"""Round-3 probe — pinpoint whether Native/Ours 挑的是同一批图.""" +import os, sys, inspect, json +import numpy as np +from PIL import Image + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")] +THUMB = 256 # reverse-lookup 降采样尺寸 + + +def sec(title): + print("\n" + "=" * 70) + print(f" {title}") + print("=" * 70) + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + args = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: args.append(source_path) + elif i == 1: args.append(img_dir) + elif k == "eval": args.append(True) + elif k == "train_test_exp": args.append(False) + else: args.append(p.default if p.default != inspect.Parameter.empty else "") + return args + + +def load_thumb(path): + img = Image.open(path).convert('RGB').resize((THUMB, THUMB), Image.LANCZOS) + return np.asarray(img, dtype=np.float32) / 255.0 + + +# ---------- Probe 11: Reverse lookup Native GT → images_2 ---------- +def probe_reverse_lookup(scene): + sec(f"PROBE 11 — Reverse-lookup Native GT in full images_2 [{scene}]") + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + gt_dir = os.path.join(cell, "gt_test_30000") + img_dir = os.path.join(DATASET_ROOT, scene, "images_2") + + native_gts = sorted([f for f in os.listdir(gt_dir) if f.lower().endswith(('.png','.jpg'))]) + i2_files = sorted([f for f in os.listdir(img_dir) if f.lower().endswith(('.png','.jpg','.jpeg'))]) + print(f"Native GT count : {len(native_gts)}") + print(f"images_2 count : {len(i2_files)}") + print(f"Building {THUMB}x{THUMB} thumb cache for vectorized PSNR match...") + + i2_arr = np.stack([load_thumb(os.path.join(img_dir, f)) for f in i2_files]) # (N, H, W, 3) + print(f" cache shape: {i2_arr.shape}") + + print(f"\n{'idx':>3} {'native_gt':>12} -> {'best_match':>15} {'best_psnr':>10} {'margin':>8}") + print("-" * 60) + matches = [] + for idx, ng in enumerate(native_gts): + ng_arr = load_thumb(os.path.join(gt_dir, ng)) + diff = i2_arr - ng_arr[None] + mse = (diff * diff).mean(axis=(1, 2, 3)) + psnrs = 10 * np.log10(1.0 / np.maximum(mse, 1e-10)) + order = np.argsort(-psnrs) + best_idx, best_psnr = order[0], psnrs[order[0]] + margin = best_psnr - psnrs[order[1]] + best_name = i2_files[best_idx] + matches.append((ng, best_name, float(best_psnr), float(margin))) + if idx < 10: + flag = "" if best_psnr > 30 else " [LOW]" + print(f"{idx:>3} {ng:>12} -> {best_name:>15} {best_psnr:>10.2f} {margin:>8.2f}{flag}") + + if len(matches) > 10: + print(f"... (showing first 10 of {len(matches)})") + + match_names = [m[1] for m in matches] + match_set = set(match_names) + low_quality = [m for m in matches if m[2] < 30] + + print(f"\n--- Reverse-lookup summary ---") + print(f"Unique matches : {len(match_set)} / {len(matches)} Native GT") + print(f"Low-quality (<30 dB) matches: {len(low_quality)} " + f"→ {'[WARN] Native GT 可能来自一个 Ours 看不到的图源' if low_quality else '[OK]'}") + + # 对比 Ours split + from scene.dataset_readers import readColmapSceneInfo + scene_info = readColmapSceneInfo(*build_args( + os.path.join(DATASET_ROOT, scene), "images_2")) + ours_set = set(c.image_name for c in scene_info.test_cameras) + + overlap = match_set & ours_set + n_only = match_set - ours_set + o_only = ours_set - match_set + print(f"\n--- Compare to Ours eval=True test split ---") + print(f"Native真挑 ∩ Ours挑 : {len(overlap):>3}") + print(f"Native only : {len(n_only):>3} e.g. {sorted(n_only)[:5]}") + print(f"Ours only : {len(o_only):>3} e.g. {sorted(o_only)[:5]}") + + if not o_only and not n_only: + print(f"\n[VERDICT] 两边挑的是完全同一批图 → 纯粹 enumerate 顺序不同") + # 检查是否 ordering 差异 + our_test_names = [c.image_name for c in scene_info.test_cameras] + if list(match_names) == our_test_names: + print(f" 且顺序也完全一致 — baseline 差异不来自 GT 端") + else: + print(f" 但 enumerate 顺序不同 — 会影响 render↔GT 配对:") + print(f" Native order (first 5): {match_names[:5]}") + print(f" Ours order (first 5): {our_test_names[:5]}") + elif len(overlap) == 0: + print(f"\n[VERDICT] 两边完全不相交 → Native 训练时的 dataset_readers 和我们用的不是同一个版本") + else: + pct = 100 * len(overlap) / max(len(ours_set), 1) + print(f"\n[VERDICT] 部分重合 ({pct:.0f}%) → split 逻辑在某维度上漂移") + + return matches, scene_info + + +# ---------- Probe 12: Baseline JSON 内容 ---------- +def probe_baseline_json(scene): + sec(f"PROBE 12 — Native baseline JSON content [{scene}]") + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + for name in ["metrics_test_iter30000.json", "results.json", + "cfg_args"]: + p = os.path.join(cell, name) + if os.path.exists(p): + print(f"\n <{name}>") + with open(p) as f: + content = f.read() + print(content[:800]) + if len(content) > 800: + print(f" ... ({len(content)} bytes total)") + + +# ---------- Probe 13: enumerate 顺序与 LLFF hold 对齐 ---------- +def probe_enumerate_order(scene, matches, scene_info): + sec(f"PROBE 13 — enumerate alignment: Native GT idx i <-> which images_2 file [{scene}]") + our_names = [c.image_name for c in scene_info.test_cameras] + print(f"{'idx':>3} {'Native→真图':>18} {'Ours test_cam[i]':>18} {'same?':>6}") + print("-" * 55) + n_align = 0 + for i, (ng, bn, bp, _) in enumerate(matches[:min(15, len(matches))]): + on = our_names[i] if i < len(our_names) else "—" + same = (bn == on) + if same: n_align += 1 + print(f"{i:>3} {bn:>18} {on:>18} {'Y' if same else 'N':>6}") + # full sweep + full_align = sum(1 for i, (_, bn, _, _) in enumerate(matches) + if i < len(our_names) and bn == our_names[i]) + print(f"\nFull enumerate alignment: {full_align} / {len(matches)} positions match") + if full_align == len(matches): + print("[OK] Native 和 Ours 的 enumerate 顺序完全一致 — 问题不在 GT/split 层") + else: + print(f"[!!] Enumerate 顺序不同 — 我们的 render[i] 对的是 Native GT[j]，i≠j") + + +def main(): + for scene, label in SCENES: + print(f"\n\n{'#'*70}\n# SCENE: {scene} [{label}]\n{'#'*70}") + matches, scene_info = probe_reverse_lookup(scene) + probe_baseline_json(scene) + probe_enumerate_order(scene, matches, scene_info) + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/probe_tnt4.py b/scripts/phase1_validation/probe_tnt4.py new file mode 100644 index 0000000000000000000000000000000000000000..c266941dc17ca99c830b52cfd8a44247f5068c9f --- /dev/null +++ b/scripts/phase1_validation/probe_tnt4.py @@ -0,0 +1,236 @@ +#!/usr/bin/env python +"""Round-4 — renderer-level diagnosis.""" +import os, sys, inspect, json, math, struct +import numpy as np +import torch +from PIL import Image +from plyfile import PlyData + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') +import gsplat + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")] + + +def sec(t): + print("\n" + "=" * 70); print(f" {t}"); print("=" * 70) + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + a = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: a.append(source_path) + elif i == 1: a.append(img_dir) + elif k == "eval": a.append(True) + elif k == "train_test_exp": a.append(False) + else: a.append(p.default if p.default != inspect.Parameter.empty else "") + return a + + +# ---------- Probe 14: cfg_args (训练超参) ---------- +def probe_cfg_args(scene): + sec(f"PROBE 14 — Training hyperparameters (cfg_args) [{scene}]") + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + for name in ["cfg_args", "args.json", "config.json", "training_args.json", + "hparams.json"]: + p = os.path.join(cell, name) + if os.path.exists(p): + print(f" <{name}>") + content = open(p).read() + print(content) + for key in ['white_background', 'sh_degree', 'resolution', 'background', + 'convert_SHs_python', 'images']: + if key in content: + print(f" [!] contains key: '{key}'") + return + print(" (no cfg_args-like file found)") + print(" listing cell dir contents:") + for f in sorted(os.listdir(cell))[:30]: + p = os.path.join(cell, f) + t = "DIR" if os.path.isdir(p) else "FIL" + print(f" {t} {f}") + + +# ---------- Probe 15: PLY SH 完整性 ---------- +def probe_ply_struct(scene): + sec(f"PROBE 15 — PLY structure & SH completeness [{scene}]") + ply = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}", + "point_cloud", "iteration_30000", "point_cloud.ply") + pd = PlyData.read(ply) + v = pd['vertex'] + props = v.data.dtype.names + print(f" N gaussians: {len(v)}") + print(f" N properties: {len(props)}") + f_rest_count = sum(1 for p in props if p.startswith('f_rest_')) + print(f" f_dc_* count: {sum(1 for p in props if p.startswith('f_dc_'))}") + print(f" f_rest_* count: {f_rest_count} " + f"(SH full 3-band needs 45; implied sh_degree=" + f"{'3' if f_rest_count == 45 else '?'})") + # opacity/scale/rot 分布 + for k in ['opacity', 'scale_0', 'f_dc_0', 'f_rest_0']: + if k in props: + vals = np.asarray(v[k]) + print(f" {k:<10} min={vals.min():.3f} max={vals.max():.3f} " + f"mean={vals.mean():.3f} std={vals.std():.3f}") + + +# ---------- Probe 16: Per-image PSNR breakdown ---------- +def probe_per_image(scene): + """直接复刻 render_single.py 的渲染循环，但这次按 Ours order 逐张 + 输出 PSNR + 和对应 Native render 同 cam 的 PSNR，找到崩最狠的几张。""" + sec(f"PROBE 16 — Per-image PSNR breakdown & delta per-cam [{scene}]") + from scene.dataset_readers import readColmapSceneInfo + from utils.graphics_utils import getWorld2View2 + from utils.general_utils import PILtoTorch + from PIL import Image, ImageOps + + source_path = os.path.join(DATASET_ROOT, scene) + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + img_dir = "images_2" + resolution = 2 + + scene_info = readColmapSceneInfo(*build_args(source_path, img_dir)) + test_cams = scene_info.test_cameras + print(f" test cams: {len(test_cams)}") + + # 构建 Native renders index: native_render_name → image + native_render_dir = os.path.join(cell, "renders_test_30000") + if os.path.isdir(os.path.join(native_render_dir, "renders")): + native_render_dir = os.path.join(native_render_dir, "renders") + native_renders = sorted([f for f in os.listdir(native_render_dir) + if f.lower().endswith(('.png', '.jpg'))]) + native_gt_dir = os.path.join(cell, "gt_test_30000") + native_gts = sorted([f for f in os.listdir(native_gt_dir) + if f.lower().endswith(('.png', '.jpg'))]) + + # 用 reverse-lookup (thumb) 找出 Native idx i 对应的真图名，从而把 + # Native render 按真图名索引 + THUMB = 128 + i2_files = sorted([f for f in os.listdir(os.path.join(source_path, img_dir)) + if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) + i2_thumbs = np.stack([ + np.asarray(Image.open(os.path.join(source_path, img_dir, f)).convert('RGB') + .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0 + for f in i2_files]) + + native_name_for = {} # native_idx → real image_name + for idx, ng in enumerate(native_gts): + t = np.asarray(Image.open(os.path.join(native_gt_dir, ng)).convert('RGB') + .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0 + mse = ((i2_thumbs - t) ** 2).mean(axis=(1, 2, 3)) + native_name_for[idx] = i2_files[int(np.argmin(mse))] + + real_to_native_idx = {v: k for k, v in native_name_for.items()} + + # 加载 PLY + pd = PlyData.read(os.path.join(cell, "point_cloud", "iteration_30000", + "point_cloud.ply")) + vv = pd['vertex'] + device = torch.device("cuda") + + def t32(x): return torch.tensor(x, dtype=torch.float32, device=device) + means = t32(np.stack((vv['x'], vv['y'], vv['z']), -1)) + quats = t32(np.stack((vv['rot_0'], vv['rot_1'], vv['rot_2'], vv['rot_3']), -1)) + scales = torch.exp(t32(np.stack((vv['scale_0'], vv['scale_1'], vv['scale_2']), -1))) + opacities = torch.sigmoid(t32(np.asarray(vv['opacity']))) + f_dc = t32(np.stack((vv['f_dc_0'], vv['f_dc_1'], vv['f_dc_2']), -1)).unsqueeze(1) + f_rest = t32(np.stack([vv[f'f_rest_{i}'] for i in range(45)], -1)) + f_rest = f_rest.view(-1, 3, 15).transpose(1, 2) + shs = torch.cat([f_dc, f_rest], dim=1) + + bg_color = torch.tensor([0., 0., 0.], device=device) + + print(f"\n {'i':>3} {'cam_name':>15} {'ours_psnr':>10} {'nat_psnr':>10} " + f"{'delta':>8} {'render_vs_render':>18}") + print(" " + "-" * 75) + + ours_psnrs, native_psnrs_this = [], [] + render_vs_render = [] + for i, c in enumerate(test_cams): + w = int(round(c.width / resolution)) + h = int(round(c.height / resolution)) + viewmat = t32(getWorld2View2(np.array(c.R), np.array(c.T))).unsqueeze(0) + fx = w / (2 * math.tan(c.FovX / 2)) + fy = h / (2 * math.tan(c.FovY / 2)) + K = t32(np.array([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]])).unsqueeze(0) + + with torch.no_grad(): + colors, _, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, + colors=shs, viewmats=viewmat, Ks=K, width=w, height=h, + sh_degree=3, packed=True, render_mode='RGB', + backgrounds=bg_color.unsqueeze(0)) + ours_rgb = colors[0].clamp(0, 1) + + # 读对应 GT + pil = Image.open(os.path.join(source_path, img_dir, c.image_name)) + pil = ImageOps.exif_transpose(pil) + gt_chw = PILtoTorch(pil, (w, h)).to(device) + if gt_chw.shape[0] == 4: gt_chw = gt_chw[:3] + gt_rgb = gt_chw.permute(1, 2, 0) + + mse = ((ours_rgb - gt_rgb) ** 2).mean().item() + ours_p = 10 * math.log10(1.0 / max(mse, 1e-10)) + ours_psnrs.append(ours_p) + + # 对应的 Native render + nat_p_str, vs_str = "—", "—" + if c.image_name in real_to_native_idx: + nidx = real_to_native_idx[c.image_name] + nr_path = os.path.join(native_render_dir, native_renders[nidx]) + ng_path = os.path.join(native_gt_dir, native_gts[nidx]) + nr = np.asarray(Image.open(nr_path).convert('RGB'), dtype=np.float32) / 255.0 + ng = np.asarray(Image.open(ng_path).convert('RGB'), dtype=np.float32) / 255.0 + if nr.shape == ng.shape: + nmse = ((nr - ng) ** 2).mean() + nat_p = 10 * math.log10(1.0 / max(nmse, 1e-10)) + native_psnrs_this.append(nat_p) + nat_p_str = f"{nat_p:.2f}" + # render vs render + ours_np = (ours_rgb.cpu().numpy() * 255).clip(0, 255).astype(np.uint8) + nr_u8 = (nr * 255).astype(np.uint8) + if ours_np.shape == nr_u8.shape: + vs_mse = ((ours_np.astype(np.float32)/255 - nr_u8.astype(np.float32)/255) ** 2).mean() + vs = 10 * math.log10(1.0 / max(vs_mse, 1e-10)) + render_vs_render.append(vs) + vs_str = f"{vs:.2f}" + + d = ours_p - (native_psnrs_this[-1] if native_psnrs_this else ours_p) + if i < 10 or (nat_p_str != "—" and abs(ours_p - float(nat_p_str)) > 2): + print(f" {i:>3} {c.image_name:>15} {ours_p:>10.2f} {nat_p_str:>10} " + f"{d:>+8.2f} {vs_str:>18}") + + print(f"\n Ours mean PSNR : {np.mean(ours_psnrs):.4f} dB") + if native_psnrs_this: + print(f" Native mean PSNR (recomputed) : {np.mean(native_psnrs_this):.4f} dB") + print(f" Mean delta : " + f"{np.mean(ours_psnrs) - np.mean(native_psnrs_this):+.4f} dB") + if render_vs_render: + arr = np.array(render_vs_render) + print(f"\n gsplat_render vs Native_render (per-cam PSNR):") + print(f" mean={arr.mean():.2f} min={arr.min():.2f} max={arr.max():.2f}") + print(f" >35 dB (near identical): {(arr>35).sum()}/{len(arr)}") + print(f" 25-35 dB : {((arr>=25)&(arr<=35)).sum()}/{len(arr)}") + print(f" <25 dB (visible diff) : {(arr<25).sum()}/{len(arr)}") + if arr.mean() > 35: + print(" [结论] 两个 renderer 输出几乎一致 → 锅不在 renderer") + elif arr.mean() > 25: + print(" [结论] renderer 有明显差异，但在合理范围") + else: + print(" [结论] renderer 输出差异巨大 → gsplat vs diff-gaussian-rasterization 不等价") + + +def main(): + for scene, label in SCENES: + print(f"\n\n{'#'*70}\n# SCENE: {scene} [{label}]\n{'#'*70}") + probe_cfg_args(scene) + probe_ply_struct(scene) + probe_per_image(scene) + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/probe_tnt5.py b/scripts/phase1_validation/probe_tnt5.py new file mode 100644 index 0000000000000000000000000000000000000000..db209c98de4977c57c3ad36771c08939433d0adc --- /dev/null +++ b/scripts/phase1_validation/probe_tnt5.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python +"""Round-5 — exposure compensation hypothesis (decisive).""" +import os, sys, json, math, inspect +import numpy as np +import torch +from PIL import Image, ImageOps +from plyfile import PlyData + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') +import gsplat + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +SCENES = [("truck", "PASS"), ("lighthouse", "FAIL")] + + +def sec(t): + print("\n" + "=" * 70); print(f" {t}"); print("=" * 70) + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + a = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: a.append(source_path) + elif i == 1: a.append(img_dir) + elif k == "eval": a.append(True) + elif k == "train_test_exp": a.append(False) + else: a.append(p.default if p.default != inspect.Parameter.empty else "") + return a + + +# ---------- Probe 17: exposure.json 结构 ---------- +def probe_exposure_structure(scene): + sec(f"PROBE 17 — exposure.json structure [{scene}]") + p = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}", "exposure.json") + if not os.path.exists(p): + print(f"[!] not found") + return None + data = json.load(open(p)) + if not isinstance(data, dict): + print(f"Not a dict: {type(data)}") + return None + keys = list(data.keys()) + print(f"Entries : {len(keys)}") + print(f"Sample keys : {keys[:5]}") + + sample = np.asarray(data[keys[0]]) + print(f"Sample shape: {sample.shape}") + print(f"Sample[0]:\n{sample}") + + all_m = np.stack([np.asarray(data[k]) for k in keys]) + print(f"\nAll exposures shape: {all_m.shape}") + + if all_m.ndim == 3 and all_m.shape[1] == 3 and all_m.shape[2] >= 3: + linear = all_m[:, :3, :3] + bias = all_m[:, :3, 3] if all_m.shape[2] == 4 else np.zeros((len(keys), 3)) + + eye = np.eye(3)[None] + lin_dev = np.abs(linear - eye).reshape(len(keys), -1).max(axis=1) + bias_mag = np.linalg.norm(bias, axis=1) + + print(f"\nLinear max |dev from I|:") + print(f" min={lin_dev.min():.5f} p50={np.percentile(lin_dev,50):.5f} " + f"p95={np.percentile(lin_dev,95):.5f} max={lin_dev.max():.5f}") + print(f"Bias magnitude ||b||:") + print(f" min={bias_mag.min():.5f} p50={np.percentile(bias_mag,50):.5f} " + f"p95={np.percentile(bias_mag,95):.5f} max={bias_mag.max():.5f}") + + trivial = ((lin_dev < 0.01) & (bias_mag < 0.01)).sum() + large = ((lin_dev >= 0.05) | (bias_mag >= 0.05)).sum() + print(f"\nNear-identity (lin_dev<0.01 & bias<0.01): {trivial} / {len(keys)}") + print(f"Large deviation (>=0.05) : {large} / {len(keys)}") + if trivial > len(keys) * 0.95: + print("[VERDICT] 曝光几乎都是单位阵 → 曝光补偿不是根因") + elif large > 0: + print("[VERDICT] 存在显著曝光偏移 → 曝光补偿极可能是根因") + return data + + +# ---------- Probe 18: 应用 exposure 后重算 PSNR ---------- +def probe_exposure_corrected(scene, exposures): + sec(f"PROBE 18 — Apply exposure to render & recompute PSNR [{scene}]") + if exposures is None: + print("[!] no exposure data, skip") + return + + from scene.dataset_readers import readColmapSceneInfo + from utils.graphics_utils import getWorld2View2 + from utils.general_utils import PILtoTorch + + source_path = os.path.join(DATASET_ROOT, scene) + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + img_dir = "images_2"; resolution = 2 + scene_info = readColmapSceneInfo(*build_args(source_path, img_dir)) + test_cams = scene_info.test_cameras + + pd = PlyData.read(os.path.join(cell, "point_cloud", "iteration_30000", "point_cloud.ply")) + v = pd['vertex'] + device = torch.device("cuda") + def t32(x): return torch.tensor(x, dtype=torch.float32, device=device) + means = t32(np.stack((v['x'], v['y'], v['z']), -1)) + quats = t32(np.stack((v['rot_0'], v['rot_1'], v['rot_2'], v['rot_3']), -1)) + scales = torch.exp(t32(np.stack((v['scale_0'], v['scale_1'], v['scale_2']), -1))) + opacities = torch.sigmoid(t32(np.asarray(v['opacity']))) + f_dc = t32(np.stack((v['f_dc_0'], v['f_dc_1'], v['f_dc_2']), -1)).unsqueeze(1) + f_rest = t32(np.stack([v[f'f_rest_{i}'] for i in range(45)], -1)).view(-1, 3, 15).transpose(1, 2) + shs = torch.cat([f_dc, f_rest], dim=1) + bg_color = torch.tensor([0., 0., 0.], device=device) + + # exposure name 映射策略：尝试多种 key 形式 + def get_exp(image_name): + candidates = [image_name, + os.path.splitext(image_name)[0], + image_name.lower(), + os.path.splitext(image_name)[0].lower()] + for c in candidates: + if c in exposures: + return torch.tensor(np.asarray(exposures[c]), dtype=torch.float32, device=device) + return None + + # sanity: check key匹配率 + matched = sum(1 for c in test_cams if get_exp(c.image_name) is not None) + print(f"Exposure key-match rate: {matched} / {len(test_cams)} test cams") + if matched == 0: + print(f" test cam names: {[c.image_name for c in test_cams[:3]]}") + print(f" exposure keys : {list(exposures.keys())[:3]}") + print("[!] 无法匹配，请检查 key 命名") + return + + raw_psnrs, corr_psnrs = [], [] + print(f"\n {'i':>3} {'cam':>14} {'raw':>8} {'corrected':>10} {'gain':>8}") + print(" " + "-" * 50) + for i, c in enumerate(test_cams): + w = int(round(c.width / resolution)); h = int(round(c.height / resolution)) + viewmat = t32(getWorld2View2(np.array(c.R), np.array(c.T))).unsqueeze(0) + fx = w / (2 * math.tan(c.FovX / 2)); fy = h / (2 * math.tan(c.FovY / 2)) + K = t32(np.array([[fx,0,w/2],[0,fy,h/2],[0,0,1]])).unsqueeze(0) + + with torch.no_grad(): + colors, _, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, + colors=shs, viewmats=viewmat, Ks=K, width=w, height=h, + sh_degree=3, packed=True, render_mode='RGB', + backgrounds=bg_color.unsqueeze(0)) + render = colors[0].clamp(0, 1) # (H, W, 3) + + pil = Image.open(os.path.join(source_path, img_dir, c.image_name)) + pil = ImageOps.exif_transpose(pil) + gt_chw = PILtoTorch(pil, (w, h)).to(device) + if gt_chw.shape[0] == 4: gt_chw = gt_chw[:3] + gt = gt_chw.permute(1, 2, 0) # (H, W, 3) + + mse_raw = ((render - gt) ** 2).mean().item() + raw_p = 10 * math.log10(1.0 / max(mse_raw, 1e-10)) + raw_psnrs.append(raw_p) + + exp_mat = get_exp(c.image_name) + if exp_mat is not None: + A = exp_mat[:3, :3]; b = exp_mat[:3, 3] if exp_mat.shape[1] >= 4 else torch.zeros(3, device=device) + # render (H,W,3) @ A.T + b (3DGS 官方公式) + corrected = render @ A.T + b + corrected = corrected.clamp(0, 1) + mse_c = ((corrected - gt) ** 2).mean().item() + corr_p = 10 * math.log10(1.0 / max(mse_c, 1e-10)) + else: + corr_p = raw_p + corr_psnrs.append(corr_p) + + if i < 15 or (corr_p - raw_p) > 2: + print(f" {i:>3} {c.image_name:>14} {raw_p:>8.2f} {corr_p:>10.2f} {corr_p - raw_p:>+8.2f}") + + raw_arr, corr_arr = np.array(raw_psnrs), np.array(corr_psnrs) + print(f"\n Ours RAW mean PSNR: {raw_arr.mean():.4f} dB") + print(f" Ours CORRECTED mean PSNR: {corr_arr.mean():.4f} dB") + + native_psnr = None + bp = os.path.join(cell, "metrics_test_iter30000.json") + if os.path.exists(bp): + bd = json.load(open(bp)) + native_psnr = bd.get("photometric", {}).get("PSNR", bd.get("PSNR")) + print(f" Native baseline PSNR : {native_psnr:.4f} dB") + if native_psnr: + print(f" Δ raw : {raw_arr.mean() - native_psnr:+.4f} dB") + print(f" Δ corrected : {corr_arr.mean() - native_psnr:+.4f} dB") + if abs(corr_arr.mean() - native_psnr) < 0.1: + print(f"\n [VERDICT CONFIRMED] 曝光补偿就是根因。") + elif abs(corr_arr.mean() - native_psnr) < abs(raw_arr.mean() - native_psnr) - 0.3: + print(f"\n [PARTIAL] 曝光补偿解释了大部分差异，但残差>0.3dB，可能还有次级因素") + else: + print(f"\n [REJECTED] 曝光补偿影响有限，需另寻根因") + + +def main(): + for scene, label in SCENES: + print(f"\n\n{'#'*70}\n# SCENE: {scene} [{label}]\n{'#'*70}") + exp = probe_exposure_structure(scene) + probe_exposure_corrected(scene, exp) + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/probe_tnt6.py b/scripts/phase1_validation/probe_tnt6.py new file mode 100644 index 0000000000000000000000000000000000000000..9ea8c114e14024fb0e73de1ea01ce55fae6dc6df --- /dev/null +++ b/scripts/phase1_validation/probe_tnt6.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python +"""Round-6 — rasterize_mode='antialiased' should kill the white haze.""" +import os, sys, inspect, math, json +import numpy as np +import torch +from PIL import Image, ImageOps +from plyfile import PlyData + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') +import gsplat + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +DUMP_ROOT = "/root/autodl-tmp/SplatAtlas/scripts/phase1_validation/diag_output" +SCENES_TEST = [("truck", "PASS"), ("lighthouse", "FAIL")] + + +def sec(t): + print("\n" + "=" * 70); print(f" {t}"); print("=" * 70) + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + a = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: a.append(source_path) + elif i == 1: a.append(img_dir) + elif k == "eval": a.append(True) + elif k == "train_test_exp": a.append(False) + else: a.append(p.default if p.default != inspect.Parameter.empty else "") + return a + + +def load_ply_tensors(ply_path, device): + v = PlyData.read(ply_path)['vertex'] + def t32(x): return torch.tensor(x, dtype=torch.float32, device=device) + means = t32(np.stack((v['x'], v['y'], v['z']), -1)) + quats = t32(np.stack((v['rot_0'], v['rot_1'], v['rot_2'], v['rot_3']), -1)) + scales = torch.exp(t32(np.stack((v['scale_0'], v['scale_1'], v['scale_2']), -1))) + opacities = torch.sigmoid(t32(np.asarray(v['opacity']))) + f_dc = t32(np.stack((v['f_dc_0'], v['f_dc_1'], v['f_dc_2']), -1)).unsqueeze(1) + f_rest = t32(np.stack([v[f'f_rest_{i}'] for i in range(45)], -1)).view(-1, 3, 15).transpose(1, 2) + shs = torch.cat([f_dc, f_rest], dim=1) + return means, quats, scales, opacities, shs + + +def render_one(means, quats, scales, opacities, shs, cam, resolution, bg, mode, device): + from utils.graphics_utils import getWorld2View2 + w = int(round(cam.width / resolution)) + h = int(round(cam.height / resolution)) + viewmat = torch.tensor(getWorld2View2(np.array(cam.R), np.array(cam.T)), + dtype=torch.float32, device=device).unsqueeze(0) + fx = w / (2 * math.tan(cam.FovX / 2)) + fy = h / (2 * math.tan(cam.FovY / 2)) + K = torch.tensor([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]], + dtype=torch.float32, device=device).unsqueeze(0) + with torch.no_grad(): + colors, _, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, + colors=shs, viewmats=viewmat, Ks=K, width=w, height=h, + sh_degree=3, packed=True, render_mode='RGB', + backgrounds=bg.unsqueeze(0), + rasterize_mode=mode) + return colors[0].clamp(0, 1), w, h + + +def load_gt(source_path, img_dir, cam, w, h, device): + from utils.general_utils import PILtoTorch + pil = Image.open(os.path.join(source_path, img_dir, cam.image_name)) + pil = ImageOps.exif_transpose(pil) + gt_chw = PILtoTorch(pil, (w, h)).to(device) + if gt_chw.shape[0] == 4: gt_chw = gt_chw[:3] + return gt_chw.permute(1, 2, 0) + + +# ---------- Probe 19: compare 'classic' vs 'antialiased' on all test cams ---------- +def probe_modes(scene, label): + sec(f"PROBE 19 — rasterize_mode comparison [{scene}] [{label}]") + from scene.dataset_readers import readColmapSceneInfo + + source_path = os.path.join(DATASET_ROOT, scene) + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + img_dir = "images_2"; resolution = 2 + scene_info = readColmapSceneInfo(*build_args(source_path, img_dir)) + test_cams = scene_info.test_cameras + + device = torch.device("cuda") + means, quats, scales, opacities, shs = load_ply_tensors( + os.path.join(cell, "point_cloud", "iteration_30000", "point_cloud.ply"), device) + bg = torch.tensor([0., 0., 0.], device=device) + + results = {} + for mode in ["classic", "antialiased"]: + psnrs = [] + for c in test_cams: + render, w, h = render_one(means, quats, scales, opacities, shs, + c, resolution, bg, mode, device) + gt = load_gt(source_path, img_dir, c, w, h, device) + mse = ((render - gt) ** 2).mean().item() + psnrs.append(10 * math.log10(1.0 / max(mse, 1e-10))) + results[mode] = np.array(psnrs) + print(f" mode='{mode}': mean PSNR = {results[mode].mean():.4f} dB " + f"(min={results[mode].min():.2f} max={results[mode].max():.2f})") + + # Native baseline + baseline_path = os.path.join(cell, "metrics_test_iter30000.json") + bd = json.load(open(baseline_path)) + native_psnr = bd.get("photometric", {}).get("PSNR", bd.get("PSNR")) + print(f" Native baseline = {native_psnr:.4f} dB") + print(f"") + print(f" Δ classic : {results['classic'].mean() - native_psnr:+.4f} dB") + print(f" Δ antialiased : {results['antialiased'].mean() - native_psnr:+.4f} dB") + + gain = results['antialiased'].mean() - results['classic'].mean() + print(f" Gain from antialiased: {gain:+.4f} dB") + + if abs(results['antialiased'].mean() - native_psnr) < 0.15: + print(f"\n ======> 完美命中 baseline。antialiased 就是对齐 3DGS 官方的正确模式。") + elif gain > 0.3: + print(f"\n ======> antialiased 显著改善但未完全命中。还有次级因素。") + else: + print(f"\n ======> antialiased 没帮上忙。换假说。") + + return results['antialiased'], test_cams + + +# ---------- Probe 20: re-dump worst cams under antialiased mode ---------- +def probe_redump(scene, antialiased_psnrs, test_cams): + sec(f"PROBE 20 — Re-render worst cams with antialiased [{scene}]") + source_path = os.path.join(DATASET_ROOT, scene) + cell = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{scene}") + dump = os.path.join(DUMP_ROOT, f"{scene}_antialiased") + os.makedirs(dump, exist_ok=True) + for f in os.listdir(dump): + os.remove(os.path.join(dump, f)) + + # worst 3 cams + idx_sorted = np.argsort(antialiased_psnrs)[:3] + device = torch.device("cuda") + means, quats, scales, opacities, shs = load_ply_tensors( + os.path.join(cell, "point_cloud", "iteration_30000", "point_cloud.ply"), device) + bg = torch.tensor([0., 0., 0.], device=device) + + native_render_dir = os.path.join(cell, "renders_test_30000") + if os.path.isdir(os.path.join(native_render_dir, "renders")): + native_render_dir = os.path.join(native_render_dir, "renders") + native_gt_dir = os.path.join(cell, "gt_test_30000") + native_renders = sorted([f for f in os.listdir(native_render_dir) + if f.lower().endswith(('.png', '.jpg'))]) + native_gts = sorted([f for f in os.listdir(native_gt_dir) + if f.lower().endswith(('.png', '.jpg'))]) + + THUMB = 128 + i2_files = sorted([f for f in os.listdir(os.path.join(source_path, "images_2")) + if f.lower().endswith(('.png', '.jpg', '.jpeg'))]) + i2_thumbs = np.stack([ + np.asarray(Image.open(os.path.join(source_path, "images_2", f)).convert('RGB') + .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0 + for f in i2_files]) + real_to_native = {} + for nidx, ng in enumerate(native_gts): + t = np.asarray(Image.open(os.path.join(native_gt_dir, ng)).convert('RGB') + .resize((THUMB, THUMB), Image.LANCZOS), dtype=np.float32) / 255.0 + real_to_native[i2_files[int(np.argmin(((i2_thumbs - t) ** 2).mean(axis=(1,2,3))))]] = nidx + + def to_u8(a): + if a.dtype == np.uint8: return a + return np.clip(a * 255, 0, 255).astype(np.uint8) + + def hstack(arrs, gap=20): + imgs = [to_u8(a) for a in arrs] + H = imgs[0].shape[0] + W = sum(i.shape[1] for i in imgs) + gap * (len(imgs) - 1) + out = np.ones((H, W, 3), dtype=np.uint8) * 240 + x = 0 + for img in imgs: + out[:, x:x + img.shape[1]] = img + x += img.shape[1] + gap + return out + + for i in idx_sorted: + c = test_cams[i] + render, w, h = render_one(means, quats, scales, opacities, shs, + c, 2, bg, "antialiased", device) + render_np = render.cpu().numpy() + if c.image_name not in real_to_native: + continue + nidx = real_to_native[c.image_name] + nr = np.asarray(Image.open(os.path.join(native_render_dir, native_renders[nidx])).convert('RGB'), + dtype=np.float32) / 255.0 + ng = np.asarray(Image.open(os.path.join(native_gt_dir, native_gts[nidx])).convert('RGB'), + dtype=np.float32) / 255.0 + if nr.shape != render_np.shape: + th, tw = render_np.shape[:2] + nr = np.asarray(Image.fromarray(to_u8(nr)).resize((tw, th), Image.LANCZOS), dtype=np.float32) / 255.0 + ng = np.asarray(Image.fromarray(to_u8(ng)).resize((tw, th), Image.LANCZOS), dtype=np.float32) / 255.0 + vs_psnr = 10*math.log10(1.0 / max(((render_np - nr)**2).mean(), 1e-10)) + print(f" cam {c.image_name}: new render_vs_native = {vs_psnr:.2f} dB") + + base = os.path.splitext(c.image_name)[0] + diff5 = np.clip(np.abs(render_np - nr) * 5, 0, 1) + panel = hstack([ng, render_np, nr, diff5]) + Image.fromarray(panel).save(os.path.join(dump, f"{base}_AA_panel.png")) + Image.fromarray(to_u8(render_np)).save(os.path.join(dump, f"{base}_AA_ours.png")) + + print(f"\n dumped to: {dump}/") + + +def main(): + for scene, label in SCENES_TEST: + print(f"\n\n{'#'*70}\n# SCENE: {scene} [{label}]\n{'#'*70}") + aa_psnrs, test_cams = probe_modes(scene, label) + probe_redump(scene, aa_psnrs, test_cams) + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/ranking_consistency.py b/scripts/phase1_validation/ranking_consistency.py new file mode 100644 index 0000000000000000000000000000000000000000..e911cf22a58c8a330595c42106c3ae30428fc3f6 --- /dev/null +++ b/scripts/phase1_validation/ranking_consistency.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python +"""Ranking consistency probe — DSLR-FAIL scenes × golden methods. +Decides whether the ~1.2 dB systematic offset still preserves within-benchmark +method ordering. +""" +import os, sys, inspect, math, json, time +import numpy as np +import torch +from PIL import Image, ImageOps +from plyfile import PlyData + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') +import gsplat +try: + from scipy.stats import spearmanr + HAS_SCIPY = True +except ImportError: + HAS_SCIPY = False + +DATASET_ROOT = "/root/autodl-tmp/dataset/tnt" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" + +METHODS = ["vanilla_3dgs", "erankgs", "ges", "lightgaussian", "opti3dgs", + "reactgs", "steepgs", "absgs", "gaussianpro", "minisplatting", "pixelgs"] +SCENES = ["palace", "lighthouse", "francis", "temple", "auditorium"] + + +def build_args(source_path, img_dir): + from scene.dataset_readers import readColmapSceneInfo + sig = inspect.signature(readColmapSceneInfo) + a = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: a.append(source_path) + elif i == 1: a.append(img_dir) + elif k == "eval": a.append(True) + elif k == "train_test_exp": a.append(False) + else: a.append(p.default if p.default != inspect.Parameter.empty else "") + return a + + +def find_cell(method, scene): + target = f"{method}_{scene}".lower() + if not os.path.isdir(OUTPUT_ROOT): + return None + for d in os.listdir(OUTPUT_ROOT): + if d.lower() == target: + return os.path.join(OUTPUT_ROOT, d) + return None + + +def load_ply_tensors(ply_path, device): + v = PlyData.read(ply_path)['vertex'] + props = v.data.dtype.names + def t32(x): return torch.tensor(x, dtype=torch.float32, device=device) + means = t32(np.stack((v['x'], v['y'], v['z']), -1)) + quats = t32(np.stack((v['rot_0'], v['rot_1'], v['rot_2'], v['rot_3']), -1)) + scales = torch.exp(t32(np.stack((v['scale_0'], v['scale_1'], v['scale_2']), -1))) + opacities = torch.sigmoid(t32(np.asarray(v['opacity']))) + n_rest = sum(1 for p in props if p.startswith('f_rest_')) + sh_per_ch = n_rest // 3 + sh_degree_map = {0: 0, 3: 1, 8: 2, 15: 3} + sh_degree = sh_degree_map.get(sh_per_ch, min(3, int(math.sqrt(sh_per_ch + 1)) - 1)) + f_dc = t32(np.stack((v['f_dc_0'], v['f_dc_1'], v['f_dc_2']), -1)).unsqueeze(1) + if n_rest > 0: + f_rest = t32(np.stack([v[f'f_rest_{i}'] for i in range(n_rest)], -1)) + f_rest = f_rest.view(-1, 3, sh_per_ch).transpose(1, 2) + shs = torch.cat([f_dc, f_rest], dim=1) + else: + shs = f_dc + return means, quats, scales, opacities, shs, sh_degree + + +def render_cell(method, scene, scene_info, device): + from utils.graphics_utils import getWorld2View2 + from utils.general_utils import PILtoTorch + + cell = find_cell(method, scene) + if cell is None: + return None, None, "NO_CELL" + + bp = os.path.join(cell, "metrics_test_iter30000.json") + native = None + if os.path.exists(bp): + bd = json.load(open(bp)) + native = bd.get("photometric", {}).get("PSNR") + if native is None: + native = bd.get("PSNR") + + ply_path = os.path.join(cell, "point_cloud", "iteration_30000", "point_cloud.ply") + if not os.path.exists(ply_path): + return None, native, "NO_PLY" + + try: + means, quats, scales, opacities, shs, sh_degree = load_ply_tensors(ply_path, device) + except Exception as e: + return None, native, f"PLY_ERR:{type(e).__name__}" + + source_path = os.path.join(DATASET_ROOT, scene) + img_dir = "images_2" + resolution = 2 + test_cams = scene_info.test_cameras + bg = torch.tensor([0., 0., 0.], device=device) + + psnrs = [] + try: + for c in test_cams: + w = int(round(c.width / resolution)) + h = int(round(c.height / resolution)) + viewmat = torch.tensor(getWorld2View2(np.array(c.R), np.array(c.T)), + dtype=torch.float32, device=device).unsqueeze(0) + fx = w / (2 * math.tan(c.FovX / 2)) + fy = h / (2 * math.tan(c.FovY / 2)) + K = torch.tensor([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]], + dtype=torch.float32, device=device).unsqueeze(0) + with torch.no_grad(): + colors, _, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, + colors=shs, viewmats=viewmat, Ks=K, width=w, height=h, + sh_degree=sh_degree, packed=True, render_mode='RGB', + backgrounds=bg.unsqueeze(0), + rasterize_mode='classic') + render = colors[0].clamp(0, 1) + pil = Image.open(os.path.join(source_path, img_dir, c.image_name)) + pil = ImageOps.exif_transpose(pil) + gt_chw = PILtoTorch(pil, (w, h)).to(device) + if gt_chw.shape[0] == 4: gt_chw = gt_chw[:3] + gt = gt_chw.permute(1, 2, 0) + mse = ((render - gt) ** 2).mean().item() + psnrs.append(10 * math.log10(1.0 / max(mse, 1e-10))) + except Exception as e: + return None, native, f"RENDER_ERR:{type(e).__name__}" + + del means, quats, scales, opacities, shs + torch.cuda.empty_cache() + return float(np.mean(psnrs)), native, "OK" + + +def print_matrix(mat, methods, scenes, title, fmt="{:>10.3f}"): + print(f"\n {title}") + header = " " + " " * 16 + " ".join(f"{s[:10]:>10}" for s in scenes) + print(header) + for mi, method in enumerate(methods): + row = " ".join( + (fmt.format(mat[mi, si]) if not np.isnan(mat[mi, si]) else f"{'—':>10}") + for si in range(len(scenes))) + print(f" {method:>16} {row}") + + +def rank_from_psnr(arr, valid_mask): + """Higher PSNR → rank 1. NaN / invalid → NaN.""" + ranks = np.full_like(arr, np.nan, dtype=float) + valid_idx = np.where(valid_mask)[0] + if len(valid_idx) == 0: + return ranks + vals = arr[valid_idx] + order = np.argsort(-vals) + for r, o in enumerate(order): + ranks[valid_idx[o]] = r + 1 + return ranks + + +def spearman_rho(a, b): + """Naive Spearman; fallback if scipy not available.""" + if HAS_SCIPY: + rho, p = spearmanr(a, b) + return rho, p + ra = np.argsort(np.argsort(a)) + rb = np.argsort(np.argsort(b)) + mean_a, mean_b = ra.mean(), rb.mean() + num = ((ra - mean_a) * (rb - mean_b)).sum() + den = math.sqrt(((ra - mean_a) ** 2).sum() * ((rb - mean_b) ** 2).sum()) + return (num / den if den > 0 else float('nan'), float('nan')) + + +def main(): + from scene.dataset_readers import readColmapSceneInfo + device = torch.device("cuda") + + print("Preloading scene infos...") + scene_infos = {} + for sc in SCENES: + source_path = os.path.join(DATASET_ROOT, sc) + if not os.path.exists(source_path): + print(f" [!] dataset path missing: {source_path}") + continue + scene_infos[sc] = readColmapSceneInfo(*build_args(source_path, "images_2")) + print(f" [preloaded] {sc}: {len(scene_infos[sc].test_cameras)} test cams") + + ours_mat = np.full((len(METHODS), len(SCENES)), np.nan) + native_mat = np.full((len(METHODS), len(SCENES)), np.nan) + status_mat = [["-"] * len(SCENES) for _ in METHODS] + + print("\nRunning ...\n") + t_all = time.time() + for mi, method in enumerate(METHODS): + for si, scene in enumerate(SCENES): + if scene not in scene_infos: + status_mat[mi][si] = "NOSRC" + continue + t0 = time.time() + ours, native, status = render_cell(method, scene, scene_infos[scene], device) + dt = time.time() - t0 + if ours is not None: ours_mat[mi, si] = ours + if native is not None: native_mat[mi, si] = native + status_mat[mi][si] = status + delta = (ours - native) if (ours is not None and native is not None) else float('nan') + ours_s = f"{ours:7.3f}" if ours is not None else " × " + native_s = f"{native:7.3f}" if native is not None else " × " + delta_s = f"{delta:+7.3f}" if not (isinstance(delta, float) and math.isnan(delta)) else " × " + print(f" [{method:<14} {scene:<11}] ours={ours_s} native={native_s} Δ={delta_s} " + f"({dt:5.1f}s {status})") + print(f"\nTotal time: {time.time() - t_all:.1f}s") + + # === Report === + print("\n" + "=" * 90) + print(" RANKING CONSISTENCY REPORT") + print("=" * 90) + + print_matrix(ours_mat, METHODS, SCENES, "OURS PSNR (gsplat classic)", "{:>10.3f}") + print_matrix(native_mat, METHODS, SCENES, "NATIVE baseline PSNR", "{:>10.3f}") + delta_mat = ours_mat - native_mat + print_matrix(delta_mat, METHODS, SCENES, "Δ (ours - native)", "{:>+10.3f}") + + valid_delta = delta_mat[~np.isnan(delta_mat)] + valid_ours = ours_mat[~np.isnan(ours_mat)] + print(f"\n === Δ stats (N={len(valid_delta)}) ===") + print(f" mean = {valid_delta.mean():+.4f} std = {valid_delta.std():.4f}") + print(f" min = {valid_delta.min():+.4f} max = {valid_delta.max():+.4f}") + print(f" |Δ|>1dB : {(np.abs(valid_delta) > 1).sum()}/{len(valid_delta)}") + print(f" |Δ|>2dB : {(np.abs(valid_delta) > 2).sum()}/{len(valid_delta)}") + print(f" |Δ|>3dB : {(np.abs(valid_delta) > 3).sum()}/{len(valid_delta)}") + + print(f"\n === OURS absolute value distribution ===") + print(f" min = {valid_ours.min():.3f} max = {valid_ours.max():.3f}") + print(f" <10 dB : {(valid_ours < 10).sum()}/{len(valid_ours)} [CATASTROPHIC]") + print(f" <15 dB : {(valid_ours < 15).sum()}/{len(valid_ours)}") + print(f" <20 dB : {(valid_ours < 20).sum()}/{len(valid_ours)}") + + # Flag catastrophic cells + cat_cells = [] + for mi in range(len(METHODS)): + for si in range(len(SCENES)): + if not np.isnan(ours_mat[mi, si]) and ours_mat[mi, si] < 15: + cat_cells.append((METHODS[mi], SCENES[si], ours_mat[mi, si], + native_mat[mi, si], delta_mat[mi, si])) + if cat_cells: + print(f"\n === Cells with OURS < 15 dB ===") + for m, s, o, n, d in sorted(cat_cells, key=lambda x: x[2]): + print(f" {m:>16} × {s:<12} ours={o:7.3f} native={n:7.3f} Δ={d:+.3f}") + + # Per-scene ranking consistency + print(f"\n === Per-scene ranking (Spearman ρ, higher PSNR = rank 1) ===") + print(f" {'scene':>12} {'ρ':>8} {'p':>8} {'N':>3} concordant_pairs") + for si, scene in enumerate(SCENES): + valid = ~np.isnan(ours_mat[:, si]) & ~np.isnan(native_mat[:, si]) + n = valid.sum() + if n < 3: + print(f" {scene:>12} {'—':>8} {'—':>8} {n:>3}") + continue + rho, p = spearman_rho(ours_mat[valid, si], native_mat[valid, si]) + r_ours = rank_from_psnr(ours_mat[:, si], valid) + r_native = rank_from_psnr(native_mat[:, si], valid) + # 相邻排名差异 + rank_diff = np.abs(r_ours - r_native) + max_rd = np.nanmax(rank_diff) + p_str = f"{p:.4f}" if not math.isnan(p) else "n/a" + print(f" {scene:>12} {rho:>+8.4f} {p_str:>8} {n:>3} max_rank_shift={max_rd:.0f}") + + # Per-method consistency (same method across scenes) + print(f"\n === Per-method ranking across {len(SCENES)} scenes ===") + for mi, method in enumerate(METHODS): + valid = ~np.isnan(ours_mat[mi, :]) & ~np.isnan(native_mat[mi, :]) + n = valid.sum() + if n < 3: continue + rho, _ = spearman_rho(ours_mat[mi, valid], native_mat[mi, valid]) + print(f" {method:>16} ρ={rho:+.4f} (N={n})") + + # Save JSON + out_json = "/root/autodl-tmp/SplatAtlas/scripts/phase1_validation/ranking_consistency.json" + out = { + "methods": METHODS, "scenes": SCENES, + "ours_psnr": ours_mat.tolist(), + "native_psnr": native_mat.tolist(), + "delta": delta_mat.tolist(), + "status": status_mat, + } + with open(out_json, "w") as f: + json.dump(out, f, indent=2, default=str) + print(f"\n Saved: {out_json}") + + # Final verdict + print(f"\n === VERDICT ===") + ok_abs = (valid_ours >= 15).all() + ok_delta = np.abs(valid_delta).mean() < 2.0 + n_rho_high = 0 + if HAS_SCIPY or not HAS_SCIPY: + for si, scene in enumerate(SCENES): + valid = ~np.isnan(ours_mat[:, si]) & ~np.isnan(native_mat[:, si]) + if valid.sum() >= 3: + rho, _ = spearman_rho(ours_mat[valid, si], native_mat[valid, si]) + if rho > 0.8: + n_rho_high += 1 + print(f" OURS absolute ≥15dB: {'YES' if ok_abs else 'NO'}") + print(f" |Δ|mean < 2dB: {'YES' if ok_delta else 'NO'}") + print(f" Scenes with ρ>0.8: {n_rho_high}/{len(SCENES)}") + if ok_abs and n_rho_high >= len(SCENES) - 1: + print(f" → Benchmark 内部一致性良好。可以收工。") + else: + print(f" → 有异常，需要看具体哪些 cell 不稳。") + + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/render_single.py b/scripts/phase1_validation/render_single.py new file mode 100644 index 0000000000000000000000000000000000000000..ea20330fc4916a7544f67ee595659ed592fbc440 --- /dev/null +++ b/scripts/phase1_validation/render_single.py @@ -0,0 +1,155 @@ +import os +import sys +import math +import argparse +import torch +import numpy as np +from PIL import Image, ImageOps +from plyfile import PlyData +import gsplat + +sys.path.insert(0, '/root/autodl-tmp/3dgsAtlas_official') + +@torch.no_grad() +def load_ply(path): + plydata = PlyData.read(path) + v = plydata['vertex'] + + means = np.stack((v['x'], v['y'], v['z']), axis=-1) + quats = np.stack((v['rot_0'], v['rot_1'], v['rot_2'], v['rot_3']), axis=-1) + scales = np.stack((v['scale_0'], v['scale_1'], v['scale_2']), axis=-1) + opacities = v['opacity'] + + f_dc = np.stack((v['f_dc_0'], v['f_dc_1'], v['f_dc_2']), axis=-1) + f_rest = np.stack([v[f'f_rest_{i}'] for i in range(45)], axis=-1) + + device = torch.device("cuda") + means = torch.tensor(means, dtype=torch.float32, device=device) + quats = torch.tensor(quats, dtype=torch.float32, device=device) + scales = torch.exp(torch.tensor(scales, dtype=torch.float32, device=device)) + opacities = torch.sigmoid(torch.tensor(opacities, dtype=torch.float32, device=device)) + + f_dc = torch.tensor(f_dc, dtype=torch.float32, device=device).unsqueeze(1) + f_rest = torch.tensor(f_rest, dtype=torch.float32, device=device) + f_rest = f_rest.view(-1, 3, 15).transpose(1, 2) + shs = torch.cat([f_dc, f_rest], dim=1) + + return means, quats, scales, opacities, shs + +def load_test_cameras(source_path, resolution): + from scene.dataset_readers import readColmapSceneInfo, readNerfSyntheticInfo + from utils.graphics_utils import getWorld2View2 + import inspect + + parsed = [] + is_synthetic = os.path.exists(os.path.join(source_path, "transforms_test.json")) and not os.path.exists(os.path.join(source_path, "sparse")) + + if not is_synthetic: + img_dir = f"images_{resolution}" if resolution > 1 and os.path.exists(os.path.join(source_path, f"images_{resolution}")) else "images" + + sig = inspect.signature(readColmapSceneInfo) + args_list = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: args_list.append(source_path) + elif i == 1: args_list.append(img_dir) + elif k == "eval": args_list.append(True) # <--- 核心修改：让 Native 决定谁是 Test 集 + elif k == "train_test_exp": args_list.append(False) + else: args_list.append(p.default if p.default != inspect.Parameter.empty else "") + + scene_info = readColmapSceneInfo(*args_list) + test_cams = scene_info.test_cameras # 直接沿用无损的 Native Test 阵列 + else: + sig = inspect.signature(readNerfSyntheticInfo) + args_list = [] + for i, (k, p) in enumerate(sig.parameters.items()): + if i == 0: args_list.append(source_path) + elif k == "eval": args_list.append(True) + elif k == "extension": args_list.append(".png") + else: args_list.append(p.default if p.default != inspect.Parameter.empty else "") + + scene_info = readNerfSyntheticInfo(*args_list) + test_cams = scene_info.test_cameras + + for c in test_cams: + w = int(round(c.width / resolution)) + h = int(round(c.height / resolution)) + viewmat = getWorld2View2(np.array(c.R), np.array(c.T)).astype(np.float32) + fx = w / (2 * math.tan(c.FovX / 2)) + fy = h / (2 * math.tan(c.FovY / 2)) + K = np.array([[fx, 0, w/2], [0, fy, h/2], [0, 0, 1]], dtype=np.float32) + + pil_img = getattr(c, 'image', None) + if pil_img is None: + img_path = getattr(c, 'image_path', None) + if not img_path or not os.path.exists(img_path): + folder = f"images_{resolution}" if resolution > 1 and os.path.exists(os.path.join(source_path, f"images_{resolution}")) else "images" + img_path = os.path.join(source_path, folder, c.image_name) + pil_img = Image.open(img_path) + + pil_img = ImageOps.exif_transpose(pil_img) + + parsed.append({ + 'name': c.image_name, + 'pil_image': pil_img, + 'viewmat': viewmat, 'K': K, 'width': w, 'height': h, + }) + return parsed + +def compute_psnr(a, b): + mse = torch.mean((a - b) ** 2) + return 100.0 if mse == 0 else (10 * torch.log10(1.0 / mse)).item() + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ply_path", required=True) + parser.add_argument("--source_path", required=True) + parser.add_argument("--model_path", required=True) + parser.add_argument("--output_dir", required=True) + parser.add_argument("--resolution", type=int, default=2) + parser.add_argument("--bg_color", type=str, default="0,0,0") + args = parser.parse_args() + + from utils.general_utils import PILtoTorch + + device = torch.device("cuda") + means, quats, scales, opacities, shs = load_ply(args.ply_path) + cameras = load_test_cameras(args.source_path, args.resolution) + + bg_parts = [float(x) for x in args.bg_color.split(",")] + bg_color = torch.tensor(bg_parts, dtype=torch.float32, device=device) + bg_color_chw = bg_color.view(3, 1, 1) + + psnrs = [] + for i, cam in enumerate(cameras): + viewmat = torch.tensor(cam['viewmat'], device=device).unsqueeze(0) + K = torch.tensor(cam['K'], device=device).unsqueeze(0) + + colors, _, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, colors=shs, + viewmats=viewmat, Ks=K, + width=cam['width'], height=cam['height'], + sh_degree=3, packed=True, render_mode='RGB', backgrounds=bg_color.unsqueeze(0), near_plane=0.01, far_plane=100.0, + ) + render_img = colors[0].clamp(0, 1) + + gt_chw_full = PILtoTorch(cam['pil_image'], (cam['width'], cam['height'])).to(device) + + if gt_chw_full.shape[0] == 4: + alpha = gt_chw_full[3:4, ...] + rgb = gt_chw_full[:3, ...] + gt_chw = rgb * alpha + bg_color_chw * (1.0 - alpha) + else: + gt_chw = gt_chw_full[:3, ...] + + gt_img = gt_chw.permute(1, 2, 0) + + psnr_val = compute_psnr(render_img, gt_img) + psnrs.append(psnr_val) + + if not psnrs: + print("\nRENDER_SINGLE DONE — STATUS: FAIL — DELTA: nan dB") + return + + mean_psnr = float(np.mean(psnrs)) + baseline_path = os.path.join(args.model_path, "metrics_test_iter30000.json") + baseline = 32.4354 diff --git a/scripts/phase1_validation/run_batch_validation.py b/scripts/phase1_validation/run_batch_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..4eaaad4ddd976a7072c5072278b93e9757091141 --- /dev/null +++ b/scripts/phase1_validation/run_batch_validation.py @@ -0,0 +1,262 @@ +import os +import sys +import subprocess +import json +import csv +import glob +from tqdm import tqdm + +BASE_DIR = "/root/autodl-tmp/SplatAtlas" +OUT_BASE = os.path.join(BASE_DIR, "outputs") +CSV_PATH = os.path.join(BASE_DIR, "scripts/phase1_validation/validation_results_anchor.csv") +ERR_LOG = os.path.join(BASE_DIR, "scripts/phase1_validation/batch_errors.log") + +METHODS = [ + "3dgsmcmc", "analyticsplatting", "atomgs", "coadaptgs", "conegs", + "erankgs", "ges", "ghap", "gof", "gslpm", "lightgaussian", "opti3dgs", + "pgsr", "reactgs", "steepgs", "vanilla_3dgs", "absgs", "gaussianpro", + "minisplatting", "pixelgs" +] + +SCENES = { + "bonsai": {"dataset": "mip360_indoor", "res": 2, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/360/bonsai"}, + "garden": {"dataset": "mip360_outdoor", "res": 4, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/360/garden"}, + "truck": {"dataset": "tnt", "res": 1, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/tnt/truck"}, + "drjohnson": {"dataset": "db", "res": 1, "bg": "0,0,0", "source": "/root/autodl-tmp/dataset/deepblending_clean/DrJohnson"} +} + +def log_error(cell_name, phase, stdout, stderr): + with open(ERR_LOG, "a") as f: + f.write(f"\n{'='*50}\n[{phase} ERROR] Cell: {cell_name}\n") + if stdout: f.write(f"--- STDOUT (last 20 lines) ---\n" + "\n".join(stdout.strip().split("\n")[-20:]) + "\n") + if stderr: f.write(f"--- STDERR (last 20 lines) ---\n" + "\n".join(stderr.strip().split("\n")[-20:]) + "\n") + f.write(f"{'='*50}\n") + +def get_actual_cell_dir(method, scene): + target = f"{method}_{scene}".lower() + if os.path.exists(OUT_BASE): + for d in os.listdir(OUT_BASE): + if d.lower() == target: + return os.path.join(OUT_BASE, d) + return None + +def phase_a_native(method, scene_key, scene_cfg, cell_dir): + json_path = os.path.join(cell_dir, "metrics_test_iter30000.json") + if os.path.exists(json_path): + try: + with open(json_path, 'r') as f: + data = json.load(f) + if "PSNR" in data or ("photometric" in data and "PSNR" in data["photometric"]): + return True + except: + pass + + # Step 1: Render + cmd_render = [ + "python", "scripts/main_render.py", "--method", method, + "--source_path", scene_cfg["source"], "--model_path", cell_dir, + "--iteration", "30000", "--resolution", str(scene_cfg["res"]) + ] + res = subprocess.run(cmd_render, cwd=BASE_DIR, capture_output=True, text=True) + if res.returncode != 0: + log_error(f"{method}_{scene_key}", "PHASE_A_RENDER", res.stdout, res.stderr) + return False + + # Step 2: Compute Metrics + py_metric = f""" +import sys, json, os +sys.path.append("ufd_evalkit") +from photometric import compute_photometric_metrics +r_dir = "{cell_dir}/renders_test_30000/renders" +g_dir = "{cell_dir}/renders_test_30000/gt" +if not os.path.exists(r_dir): + r_dir = "{cell_dir}/renders_test_30000" + g_dir = "{cell_dir}/gt_test_30000" +metrics = compute_photometric_metrics(r_dir, g_dir) +with open("{json_path}", "w") as f: + json.dump({{"PSNR": metrics.get("PSNR", 0)}}, f, indent=4) +""" + res2 = subprocess.run(["python", "-c", py_metric], cwd=BASE_DIR, capture_output=True, text=True) + if res2.returncode != 0: + log_error(f"{method}_{scene_key}", "PHASE_A_METRIC", res2.stdout, res2.stderr) + return False + return True + +def phase_b_gsplat(ply_path, scene_cfg, cell_dir, cell_name): + out_dir = os.path.join(OUT_BASE, "phase1_validation", cell_name) + os.makedirs(out_dir, exist_ok=True) + cmd = [ + "python", "scripts/phase1_validation/render_single.py", + "--ply_path", ply_path, "--source_path", scene_cfg["source"], + "--model_path", cell_dir, "--output_dir", out_dir, + "--resolution", str(scene_cfg["res"]), "--bg_color", scene_cfg["bg"] + ] + res = subprocess.run(cmd, cwd=BASE_DIR, capture_output=True, text=True) + + if res.returncode != 0: + notes = "ERROR_B" + if "CUDA out of memory" in res.stderr or "CUDA out of memory" in res.stdout: + notes = "OOM" + elif "plyfile" in res.stderr and "format" in res.stderr: + notes = "PLY format incompatible" + log_error(cell_name, "PHASE_B_GSPLAT", res.stdout, res.stderr) + return None, None, None, notes + + native_psnr, gsplat_psnr, delta = None, None, None + for line in res.stdout.split('\n'): + if "Mean PSNR (ours/gsplat)" in line: + try: gsplat_psnr = float(line.split(":")[1].strip().split()[0]) + except: pass + if "Mean PSNR (native baseline)" in line: + try: native_psnr = float(line.split(":")[1].strip().split()[0]) + except: pass + if "Delta" in line and "dB" in line and "STATUS" not in line: + try: delta = float(line.split(":")[1].strip().split()[0]) + except: pass + + # Fallback to json if regex fails + if gsplat_psnr is None: + try: + with open(os.path.join(out_dir, "psnr_results.json"), "r") as f: + d = json.load(f) + gsplat_psnr = float(np.mean(list(d.values()))) + except: pass + + notes = "" + if gsplat_psnr is not None and gsplat_psnr < 15.0: + notes = "quantized PLY? low PSNR" + + return native_psnr, gsplat_psnr, delta, notes + +def main(): + print("=== SplatAtlas Phase 1 Task B: Anchor Batch Validation ===") + + completed = set() + write_header = not os.path.exists(CSV_PATH) + if not write_header: + with open(CSV_PATH, "r") as f: + reader = csv.DictReader(f) + for row in reader: + completed.add((row['method'], row['scene'])) + + total_cells = len(METHODS) * len(SCENES) + pbar = tqdm(total=total_cells, desc="Validating Cells", unit="cell") + + with open(CSV_PATH, "a", newline='') as f: + writer = csv.writer(f) + if write_header: + writer.writerow(["method", "scene", "dataset", "resolution", "has_ply", "native_psnr", "gsplat_psnr", "delta_db", "status", "notes"]) + + for method in METHODS: + for scene, cfg in SCENES.items(): + if (method, scene) in completed: + pbar.update(1) + continue + + cell_dir = get_actual_cell_dir(method, scene) + if not cell_dir: + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "False", "", "", "", "SKIP", "Directory missing"]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] SKIP (No Dir)") + pbar.update(1) + continue + + ply_path = os.path.join(cell_dir, "point_cloud/iteration_30000/point_cloud.ply") + if not os.path.exists(ply_path): + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "False", "", "", "", "SKIP", "PLY missing"]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] SKIP (No PLY)") + pbar.update(1) + continue + + cell_name = os.path.basename(cell_dir) + + # Phase A + if not phase_a_native(method, scene, cfg, cell_dir): + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_A", "Native render/eval failed"]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] ERROR_A") + pbar.update(1) + continue + + # Phase B + native, gsplat, delta, notes = phase_b_gsplat(ply_path, cfg, cell_dir, cell_name) + + if native is None and gsplat is None: + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_B", notes]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] ERROR_B") + else: + if delta is None and native is not None and gsplat is not None: + delta = gsplat - native + + status = "FAIL" + if delta is not None: + if abs(delta) < 0.3: status = "PASS" + elif abs(delta) < 0.5: status = "WARN" + + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", + f"{native:.4f}" if native else "", + f"{gsplat:.4f}" if gsplat else "", + f"{delta:.4f}" if delta else "", + status, notes]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] {status} Δ={delta:+.2f}dB") + + pbar.update(1) + + pbar.close() + + # Generate Summary + stats = {"PASS": 0, "WARN": 0, "FAIL": 0, "SKIP": 0, "ERROR": 0} + method_summary = {m: {"PASS": 0, "total_run": 0, "fails": []} for m in METHODS} + + with open(CSV_PATH, "r") as f: + reader = csv.DictReader(f) + for row in reader: + st = row['status'] + if st.startswith("ERROR"): stats["ERROR"] += 1 + elif st in stats: stats[st] += 1 + + m = row['method'] + if m in method_summary and st not in ("SKIP", "ERROR_A"): + method_summary[m]["total_run"] += 1 + if st == "PASS": method_summary[m]["PASS"] += 1 + elif st in ("FAIL", "WARN", "ERROR_B"): + method_summary[m]["fails"].append(f"{row['scene']} Δ={row['delta_db']} ({st})") + + total = sum(stats.values()) + print("\n" + "="*40) + print("==== BATCH SUMMARY ====") + print(f" PASS: {stats['PASS']}/{total}") + print(f" WARN: {stats['WARN']}") + print(f" FAIL: {stats['FAIL']}") + print(f" SKIP: {stats['SKIP']}") + print(f" ERROR: {stats['ERROR']}") + print("\nFAIL/WARN Details by Method:") + + all_pass = [] + all_fail = [] + mixed = [] + + for m, d in method_summary.items(): + if d["total_run"] == 0: continue + if d["PASS"] == d["total_run"]: + all_pass.append(m) + elif d["PASS"] == 0: + all_fail.append(m) + print(f" {m}: ALL FAIL/ERROR -> " + ", ".join(d["fails"])) + else: + mixed.append(m) + print(f" {m}: Mixed -> " + ", ".join(d["fails"])) + + print("\nMethod Classification:") + print(f" Fully PASS ({len(all_pass)}): " + ", ".join(all_pass)) + print(f" Fully FAIL/ERR ({len(all_fail)}) [Needs DataLoader Branch]: " + ", ".join(all_fail)) + print(f" Mixed ({len(mixed)}) [Scene-dependent issues]: " + ", ".join(mixed)) + print("========================================") + print("BATCH ANCHOR DONE") + print(f"See: {CSV_PATH}") + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/run_full_batch.py b/scripts/phase1_validation/run_full_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..f7834dcac4f29f20c186e5de61a42e851d8fe344 --- /dev/null +++ b/scripts/phase1_validation/run_full_batch.py @@ -0,0 +1,223 @@ +import os +import sys +import subprocess +import json +import csv +import glob +from tqdm import tqdm +import numpy as np + +BASE_DIR = "/root/autodl-tmp/SplatAtlas" +OUT_BASE = os.path.join(BASE_DIR, "outputs") +DATASET_ROOT = "/root/autodl-tmp/dataset" +CSV_PATH = os.path.join(BASE_DIR, "scripts/phase1_validation/validation_results_final.csv") +ERR_LOG = os.path.join(BASE_DIR, "scripts/phase1_validation/final_batch_errors.log") + +# 全量 20 种方法 +ALL_METHODS = [ + "3dgsmcmc", "analyticsplatting", "atomgs", "coadaptgs", "conegs", + "erankgs", "ges", "ghap", "gof", "gslpm", "lightgaussian", "opti3dgs", + "pgsr", "reactgs", "steepgs", "vanilla_3dgs", "absgs", "gaussianpro", + "minisplatting", "pixelgs" +] + +def build_scene_database(): + """绝对白名单机制，屏蔽衍生目录""" + db = {} + allowed_roots = ["360", "tnt", "deepblending_clean"] + + for root in allowed_roots: + search_path = os.path.join(DATASET_ROOT, root, "*") + for path in glob.glob(search_path): + if not os.path.isdir(path): continue + + scene_name = os.path.basename(path).lower() + + if root == "tnt": + res, dataset_type = 2, "tnt" + elif root == "deepblending_clean": + res, dataset_type = 1, "db" + elif root == "360": + outdoor = ['garden', 'stump', 'treehill', 'bicycle', 'flowers'] + res = 4 if any(o in scene_name for o in outdoor) else 2 + dataset_type = "mip360_outdoor" if res == 4 else "mip360_indoor" + + db[scene_name] = {"source": path, "res": res, "dataset": dataset_type, "bg": "0,0,0"} + return db + +def log_error(cell_name, phase, stdout, stderr): + with open(ERR_LOG, "a") as f: + f.write(f"\n{'='*50}\n[{phase} ERROR] Cell: {cell_name}\n") + if stdout: f.write(f"--- STDOUT ---\n" + "\n".join(stdout.strip().split("\n")[-20:]) + "\n") + if stderr: f.write(f"--- STDERR ---\n" + "\n".join(stderr.strip().split("\n")[-20:]) + "\n") + +def get_actual_cell_dir(method, scene): + target = f"{method}_{scene}".lower() + if os.path.exists(OUT_BASE): + for d in os.listdir(OUT_BASE): + if d.lower() == target: + return os.path.join(OUT_BASE, d) + return None + +def phase_a_native(method, scene_key, scene_cfg, cell_dir): + json_path = os.path.join(cell_dir, "metrics_test_iter30000.json") + if os.path.exists(json_path): + try: + with open(json_path, 'r') as f: + data = json.load(f) + if "PSNR" in data or ("photometric" in data and "PSNR" in data["photometric"]): + return True + except: pass + + cmd_render = [ + "python", "scripts/main_render.py", "--method", method, + "--source_path", scene_cfg["source"], "--model_path", cell_dir, + "--iteration", "30000", "--resolution", str(scene_cfg["res"]) + ] + res = subprocess.run(cmd_render, cwd=BASE_DIR, capture_output=True, text=True) + if res.returncode != 0: + log_error(f"{method}_{scene_key}", "PHASE_A_RENDER", res.stdout, res.stderr) + return False + + py_metric = f""" +import sys, json, os +sys.path.append("ufd_evalkit") +from photometric import compute_photometric_metrics +r_dir = "{cell_dir}/renders_test_30000/renders" +g_dir = "{cell_dir}/renders_test_30000/gt" +if not os.path.exists(r_dir): + r_dir = "{cell_dir}/renders_test_30000" + g_dir = "{cell_dir}/gt_test_30000" +metrics = compute_photometric_metrics(r_dir, g_dir) +with open("{json_path}", "w") as f: + json.dump({{"PSNR": metrics.get("PSNR", 0)}}, f, indent=4) +""" + res2 = subprocess.run(["python", "-c", py_metric], cwd=BASE_DIR, capture_output=True, text=True) + if res2.returncode != 0: + log_error(f"{method}_{scene_key}", "PHASE_A_METRIC", res2.stdout, res2.stderr) + return False + return True + +def phase_b_gsplat(ply_path, scene_cfg, cell_dir, cell_name): + out_dir = os.path.join(OUT_BASE, "phase1_validation", cell_name) + os.makedirs(out_dir, exist_ok=True) + cmd = [ + "python", "scripts/phase1_validation/render_single.py", + "--ply_path", ply_path, "--source_path", scene_cfg["source"], + "--model_path", cell_dir, "--output_dir", out_dir, + "--resolution", str(scene_cfg["res"]), "--bg_color", scene_cfg["bg"] + ] + res = subprocess.run(cmd, cwd=BASE_DIR, capture_output=True, text=True) + + if res.returncode != 0: + log_error(cell_name, "PHASE_B_GSPLAT", res.stdout, res.stderr) + return None, None, None, "ERROR_B" + + native_psnr, gsplat_psnr, delta = None, None, None + for line in res.stdout.split('\n'): + if "Mean PSNR (ours/gsplat)" in line: + try: gsplat_psnr = float(line.split(":")[1].strip().split()[0]) + except: pass + if "Mean PSNR (native baseline)" in line: + try: native_psnr = float(line.split(":")[1].strip().split()[0]) + except: pass + if "Delta" in line and "dB" in line and "STATUS" not in line: + try: delta = float(line.split(":")[1].strip().split()[0]) + except: pass + + if gsplat_psnr is None: + try: + with open(os.path.join(out_dir, "psnr_results.json"), "r") as f: + d = json.load(f) + gsplat_psnr = float(np.mean(list(d.values()))) + except: pass + + notes = "" + if gsplat_psnr is not None and gsplat_psnr < 15.0: + notes = "low PSNR" + + return native_psnr, gsplat_psnr, delta, notes + +def main(): + print("=== SplatAtlas Phase 1 Final Full Batch ===") + SCENES = build_scene_database() + print(f"[Info] Loaded {len(SCENES)} verified scenes.") + + completed = set() + write_header = not os.path.exists(CSV_PATH) + if not write_header: + with open(CSV_PATH, "r") as f: + reader = csv.DictReader(f) + for row in reader: + completed.add((row['method'], row['scene'])) + + valid_cells = [] + for method in ALL_METHODS: + for scene, cfg in SCENES.items(): + cell_dir = get_actual_cell_dir(method, scene) + if cell_dir: + valid_cells.append((method, scene, cfg, cell_dir)) + + print(f"[Info] Found {len(valid_cells)} pre-trained cells across {len(ALL_METHODS)} methods.\n") + + pbar = tqdm(total=len(valid_cells), desc="Validating Cells", unit="cell") + + with open(CSV_PATH, "a", newline='') as f: + writer = csv.writer(f) + if write_header: + writer.writerow(["method", "scene", "dataset", "resolution", "has_ply", "native_psnr", "gsplat_psnr", "delta_db", "status", "notes"]) + + for method, scene, cfg, cell_dir in valid_cells: + if (method, scene) in completed: + pbar.update(1) + continue + + ply_path = os.path.join(cell_dir, "point_cloud/iteration_30000/point_cloud.ply") + if not os.path.exists(ply_path): + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "False", "", "", "", "SKIP", "PLY missing"]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] SKIP") + pbar.update(1) + continue + + cell_name = os.path.basename(cell_dir) + + if not phase_a_native(method, scene, cfg, cell_dir): + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_A", "Native render failed"]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] ERROR_A") + pbar.update(1) + continue + + native, gsplat, delta, notes = phase_b_gsplat(ply_path, cfg, cell_dir, cell_name) + + if native is None and gsplat is None: + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", "", "", "", "ERROR_B", notes]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] ERROR_B") + else: + if delta is None and native is not None and gsplat is not None: + delta = gsplat - native + + status = "PASS" # 基于你的要求，弱化硬阈值，仅记录数据 + if delta is not None and abs(delta) >= 1.0: + status = "FAIL" + notes = f"Systematic Bias | {notes}".strip(" | ") + + writer.writerow([method, scene, cfg["dataset"], cfg["res"], "True", + f"{native:.4f}" if native else "", + f"{gsplat:.4f}" if gsplat else "", + f"{delta:.4f}" if delta else "", + status, notes]) + f.flush() + pbar.set_postfix_str(f"[{method}/{scene}] {status}") + + pbar.update(1) + + pbar.close() + print("\n========================================") + print(f"FINAL BATCH COMPLETED. Results saved to: {CSV_PATH}") + print("========================================") + +if __name__ == "__main__": + main() diff --git a/scripts/phase1_validation/sanity_check.py b/scripts/phase1_validation/sanity_check.py new file mode 100644 index 0000000000000000000000000000000000000000..dc348a457108732411bcfd496e3c17da14cd469f --- /dev/null +++ b/scripts/phase1_validation/sanity_check.py @@ -0,0 +1,45 @@ +import subprocess +import os + +print("=== Phase 1: Minimal Sanity Check (2x1) ===") + +# 强制硬编码路径，物理隔离 360_sparse_ratio8 的干扰 +TASKS = [ + { + "method": "vanilla_3dgs", + "scene": "bonsai", + "source": "/root/autodl-tmp/dataset/360/bonsai", + "model": "/root/autodl-tmp/SplatAtlas/outputs/vanilla_3dgs_bonsai", + "res": 2 + }, + { + "method": "pixelgs", + "scene": "bonsai", + "source": "/root/autodl-tmp/dataset/360/bonsai", + "model": "/root/autodl-tmp/SplatAtlas/outputs/pixelgs_bonsai", + "res": 2 + } +] + +for task in TASKS: + print(f"\n>> Testing {task['method']} on {task['scene']}...") + cmd = [ + "python", "scripts/phase1_validation/render_single.py", + "--ply_path", f"{task['model']}/point_cloud/iteration_30000/point_cloud.ply", + "--source_path", task["source"], + "--model_path", task["model"], + "--output_dir", f"/tmp/sanity_{task['method']}", + "--resolution", str(task["res"]), + "--bg_color", "0,0,0" + ] + res = subprocess.run(cmd, capture_output=True, text=True) + + # 提取关键指标 + for line in res.stdout.split("\n"): + if "Mean PSNR" in line or "Delta" in line or "STATUS" in line: + print(f" {line}") + if res.returncode != 0: + print(f" [CRITICAL] Process crashed with code {res.returncode}") + print(f" STDERR: {res.stderr[-200:]}") + +print("\n=== Sanity Check Completed ===") diff --git a/scripts/phase1_validation/sanity_tnt.py b/scripts/phase1_validation/sanity_tnt.py new file mode 100644 index 0000000000000000000000000000000000000000..ffefd308e1dfde6af51d09081255dd219a909bf4 --- /dev/null +++ b/scripts/phase1_validation/sanity_tnt.py @@ -0,0 +1,17 @@ +import subprocess +s = "lighthouse" +print(f"\n>> Testing vanilla_3dgs on {s} (EXIF/Images2 Match)...") + +cmd = [ + "python", "scripts/phase1_validation/render_single.py", + "--ply_path", f"/root/autodl-tmp/SplatAtlas/outputs/vanilla_3dgs_{s}/point_cloud/iteration_30000/point_cloud.ply", + "--source_path", f"/root/autodl-tmp/dataset/tnt/{s}", + "--model_path", f"/root/autodl-tmp/SplatAtlas/outputs/vanilla_3dgs_{s}", + "--output_dir", f"/tmp/sanity_tnt_{s}", + "--resolution", "2", + "--bg_color", "0,0,0" +] +res = subprocess.run(cmd, capture_output=True, text=True, cwd="/root/autodl-tmp/SplatAtlas") +for line in res.stdout.split("\n"): + if "Mean PSNR" in line or "Delta" in line or "STATUS" in line: + print(f" {line}") diff --git a/scripts/phase2_counterfactual/aggregate_part3a.py b/scripts/phase2_counterfactual/aggregate_part3a.py new file mode 100644 index 0000000000000000000000000000000000000000..6704096f1e12d992d9f52cf881e0b3b682d8dc7e --- /dev/null +++ b/scripts/phase2_counterfactual/aggregate_part3a.py @@ -0,0 +1,162 @@ +import os, sys, json, numpy as np +from datetime import datetime, timezone + +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility" +PROF_DIR = os.path.join(OUTPUT_ROOT, "_profiles") +SCENES = [ + "bicycle", "bonsai", "counter", "flowers", "garden", "kitchen", "room", "stump", "treehill", + "auditorium", "ballroom", "barn", "caterpillar", "courtroom", "lighthouse", "museum", "palace", "playground", "temple", "train", "truck", + "DrJohnson", "Playroom", + "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship" +] +METHODS = ["3dgsmcmc", "absgs", "gaussianpro", "pixelgs"] +ALL_METHODS = ["vanilla_3dgs"] + METHODS +METRICS = ["sh_net_effect", "sh_corruption_rate", "opacity_net_effect", "coverage", "coverage_error_fraction", "residual_error"] + +def get_iqr(arr): return float(np.percentile(arr, 75) - np.percentile(arr, 25)) if len(arr)>0 else 0.0 + +# 1. 读取 Vanilla 基准 +with open(os.path.join(PROF_DIR, "vanilla_3dgs_profile.json"), "r") as f: + vanilla_prof = json.load(f) + +profiles = {"vanilla_3dgs": vanilla_prof} +method_psnr_medians = {} + +vanilla_psnrs = [] +for s in SCENES: + p = os.path.join(OUTPUT_ROOT, f"vanilla_3dgs_{s}", "counterfactual_summary.json") + if os.path.exists(p): + with open(p) as f: vanilla_psnrs.append(json.load(f).get("psnr_full", 0)) +method_psnr_medians["vanilla_3dgs"] = float(np.median(vanilla_psnrs)) if vanilla_psnrs else 0.0 + +sh_negative_records = [] +global_pass_count = 0 +global_fail_count = 0 +total_cells = 0 + +# 2. 提取 4 种变体方法的独立 Profile +for method in METHODS: + n_pass = n_warn = n_fail = n_sh_neg = 0 + failed_scenes = [] + metrics_store = {k: [] for k in METRICS} + psnr_store = [] + ds_stats = {ds: {"n_raw": 0, "n_valid": 0, "metrics": {k: [] for k in METRICS}} for ds in ["mip360", "tnt", "db", "synthetic"]} + + for scene in SCENES: + p = os.path.join(OUTPUT_ROOT, f"{method}_{scene}", "counterfactual_summary.json") + if not os.path.exists(p): continue + total_cells += 1 + with open(p) as f: d = json.load(f) + + ds = d.get("dataset", "unknown") + if ds in ds_stats: ds_stats[ds]["n_raw"] += 1 + + st = d.get("sanity_status", "FAIL") + if st in ["PASS", "WARN"]: + if st == "PASS": n_pass += 1 + else: n_warn += 1 + global_pass_count += 1 + else: + n_fail += 1; global_fail_count += 1 + failed_scenes.append(f"{scene}: {','.join(d.get('sanity_failures', []))}") + continue + + if d.get("sh_negative_flag", False): + n_sh_neg += 1 + sh_negative_records.append(f"{method} on {scene}") + + if ds in ds_stats: ds_stats[ds]["n_valid"] += 1 + for k in METRICS: + val = d.get(k) + if val is not None: + metrics_store[k].append(val) + if ds in ds_stats: ds_stats[ds]["metrics"][k].append(val) + psnr_store.append(d.get("psnr_full", 0)) + + method_psnr_medians[method] = float(np.median(psnr_store)) if psnr_store else 0.0 + + prof = { + "method": method, "n_scenes_total": len(SCENES), + "n_scenes_pass": n_pass, "n_scenes_warn": n_warn, "n_scenes_fail": n_fail, "failed_scenes": failed_scenes, + "per_dataset": {}, "cross_scene_median": {}, "cross_scene_iqr": {}, + "generated_at": datetime.now(timezone.utc).isoformat(), + "n_scenes_sh_negative": n_sh_neg, + "sh_sign_consistency": "positive_all" if n_sh_neg == 0 else ("negative_all" if n_sh_neg == len(SCENES) else "mixed"), + "delta_vs_vanilla": {} + } + + for ds, ds_data in ds_stats.items(): + ds_dict = {"n_scenes": ds_data["n_valid"], "n_scenes_raw": ds_data["n_raw"]} + for k in METRICS: ds_dict[f"median_{k}"] = float(np.median(ds_data["metrics"][k])) if ds_data["metrics"][k] else None + if ds == "db": ds_dict["caveat"] = "n=2, statistics not robust" + if ds == "synthetic": ds_dict["caveat"] = "not in Phase 1 CSV; sanity#1 skipped" + prof["per_dataset"][ds] = ds_dict + + for k in METRICS: + m_val = float(np.median(metrics_store[k])) if metrics_store[k] else 0.0 + prof["cross_scene_median"][k] = m_val + prof["cross_scene_iqr"][k] = get_iqr(metrics_store[k]) + v_val = vanilla_prof.get("cross_scene_median", {}).get(k, 0.0) + prof["delta_vs_vanilla"][f"{k}_median"] = m_val - v_val + + with open(os.path.join(PROF_DIR, f"{method}_profile.json"), "w") as f: json.dump(prof, f, indent=4) + profiles[method] = prof + +# 3. 生成跨方法聚合对比字典 +cross_prof = { + "methods": ALL_METHODS, "metrics": {}, + "psnr_cross_method_iqr": get_iqr(list(method_psnr_medians.values())), + "counterfactual_vs_psnr_iqr_ratio": {} +} + +for k in METRICS: + medians_dict = {m: profiles[m].get("cross_scene_median", {}).get(k, 0.0) for m in ALL_METHODS} + vals = list(medians_dict.values()) + cross_prof["metrics"][k] = { + "per_method_median": medians_dict, + "cross_method_iqr": get_iqr(vals), + "cross_method_range": float(np.max(vals) - np.min(vals)) + } + +psnr_iqr = cross_prof["psnr_cross_method_iqr"] +if psnr_iqr > 0: + cross_prof["counterfactual_vs_psnr_iqr_ratio"] = { + "sh_net": cross_prof["metrics"]["sh_net_effect"]["cross_method_iqr"] / psnr_iqr, + "opa_net": cross_prof["metrics"]["opacity_net_effect"]["cross_method_iqr"] / psnr_iqr, + "coverage_err": cross_prof["metrics"]["coverage_error_fraction"]["cross_method_iqr"] / psnr_iqr, + "residual": cross_prof["metrics"]["residual_error"]["cross_method_iqr"] / psnr_iqr + } + +with open(os.path.join(PROF_DIR, "part3a_cross_method_comparison.json"), "w") as f: json.dump(cross_prof, f, indent=4) + +# 4. 生成终端/Markdown输出报告 +md = [ + "## Phase 2 Part 3a Summary", + f"- **Stage 1 Verdict**: Property (Healthy cross-scene variance, proceeded to stage 2).", + f"- **Stage 2 Verdict**: 5/5 PASS on bonsai pilot.", + f"- **Stage 3 Execution**: {total_cells} cells processed. PASS/WARN: {global_pass_count}, FAIL: {global_fail_count}.", + "", "### SH Sign Consistency", "| Method | Sign Consistency | Negative Count |", "|---|---|---|" +] +for m in METHODS: md.append(f"| {m} | {profiles[m]['sh_sign_consistency']} | {profiles[m]['n_scenes_sh_negative']} |") + +md.extend(["", "### Delta vs Vanilla (Medians)", "| Method | d_SH_net | d_Opa_net | d_Cov | d_CovErr | d_Resid | d_SH_corr |", "|---|---|---|---|---|---|---|"]) +for m in METHODS: + d = profiles[m]["delta_vs_vanilla"] + md.append(f"| {m} | {d['sh_net_effect_median']:.4f} | {d['opacity_net_effect_median']:.4f} | {d['coverage_median']:.4f} | {d['coverage_error_fraction_median']:.4f} | {d['residual_error_median']:.4f} | {d['sh_corruption_rate_median']:.4f} |") + +md.extend([ + "", "### Counterfactual vs PSNR IQR Ratio", + f"- PSNR Cross-Method IQR: **{psnr_iqr:.4f}**", + f"- SH Net Ratio: **{cross_prof['counterfactual_vs_psnr_iqr_ratio'].get('sh_net', 0):.2f}**", + f"- Opacity Net Ratio: **{cross_prof['counterfactual_vs_psnr_iqr_ratio'].get('opa_net', 0):.2f}**", + f"- Coverage Error Ratio: **{cross_prof['counterfactual_vs_psnr_iqr_ratio'].get('coverage_err', 0):.2f}**", + f"- Residual Error Ratio: **{cross_prof['counterfactual_vs_psnr_iqr_ratio'].get('residual', 0):.2f}**", + "> *Interpretation*: Ratios > 2 indicate the counterfactual metric possesses significantly higher method-level discriminative power than standard PSNR.", + "", "### Anomalies & Part 3b Readiness", + f"- SH Negative Incidents: {len(sh_negative_records)} ({', '.join(sh_negative_records[:5]) + '...' if len(sh_negative_records)>5 else 'None'}).", + "- **Signal Note**: 3dgsmcmc SH corruption rate divergence remains a key phenomenon (to be isolated in analysis).", + "- **Part 3b Brief Check**: The infrastructure safely handled 124 cells. Delta bounds for cov_err_frac were appropriately relaxed. No hard blockers for Part 3b (9-method scaling)." +]) + +print("\n" + "\n".join(md) + "\n") +with open(os.path.join(OUTPUT_ROOT, "part3a_summary.md"), "w") as f: f.write("\n".join(md)) diff --git a/scripts/phase2_counterfactual/aggregate_part3b.py b/scripts/phase2_counterfactual/aggregate_part3b.py new file mode 100644 index 0000000000000000000000000000000000000000..a94da2d72abb098b4acf92726c205264221e2a87 --- /dev/null +++ b/scripts/phase2_counterfactual/aggregate_part3b.py @@ -0,0 +1,75 @@ +import os, json, numpy as np +from datetime import datetime, timezone + +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility" +PROF_DIR = os.path.join(OUTPUT_ROOT, "_profiles") + +METHODS = ['vanilla_3dgs', 'analyticsplatting', 'erankgs', 'ges', 'lightgaussian', 'minisplatting', 'opti3dgs', 'pgsr', 'steepgs'] +SCENES = ['bicycle','bonsai','counter','flowers','garden','kitchen','room','stump','treehill', + 'auditorium','ballroom','barn','caterpillar','courtroom','lighthouse','museum','palace','playground','temple','train','truck', + 'DrJohnson','Playroom', + 'Chair','Drums','Ficus','Hotdog','Lego','Materials','Mic','Ship'] + +METRICS = ["psnr_full", "sh_net_effect", "sh_corruption_rate", "opacity_net_effect", "coverage", "coverage_error_fraction", "residual_error"] + +# 1. 抓取所有 9 个方法在 31 个场景上的中位数 (Cross-Scene Median) +method_medians = {m: {} for m in METHODS} +for m in METHODS: + for metric in METRICS: + vals = [] + for s in SCENES: + p = os.path.join(OUTPUT_ROOT, f"{m}_{s}", "counterfactual_summary.json") + if os.path.exists(p): + with open(p) as f: + d = json.load(f) + if d.get(metric) is not None: + vals.append(d[metric]) + method_medians[m][metric] = float(np.median(vals)) if vals else 0.0 + +# 2. 计算相对 IQR 统计量 (跨 9 种方法) +def calc_iqr_stats(vals_list): + if not vals_list: return None, None, None + med = float(np.median(vals_list)) + abs_iqr = float(np.percentile(vals_list, 75) - np.percentile(vals_list, 25)) + rel_iqr = (abs_iqr / abs(med) * 100) if med != 0 else None + return med, abs_iqr, rel_iqr + +cross_method_analysis = {"n_methods": len(METHODS), "methods_included": METHODS, "metrics": {}} + +# 提取 PSNR 基础倍率 +psnr_vals = [method_medians[m]["psnr_full"] for m in METHODS] +_, _, psnr_rel_iqr = calc_iqr_stats(psnr_vals) + +for metric in METRICS: + vals = [method_medians[m][metric] for m in METHODS] + med, abs_iqr, rel_iqr = calc_iqr_stats(vals) + ratio_rel = (rel_iqr / psnr_rel_iqr) if (rel_iqr is not None and psnr_rel_iqr) else None + + cross_method_analysis["metrics"][metric] = { + "median_of_9_methods": med, + "relative_iqr_pct": rel_iqr, + "ratio_vs_psnr_relative": ratio_rel + } + +# 3. 输出报告 +md_lines = [ + "# Phase 2 Part 3b — Final Heterogeneous Array Findings", + f"\n**Scope**: {len(METHODS)} methods (Vanilla + 8 heterogeneous), {len(SCENES)} scenes.", + "\n## Ultimate Discriminative Power (Relative IQR)", + f"> **Baseline**: PSNR Relative IQR across 9 methods = **{psnr_rel_iqr:.2f}%**", + "\n| Metric | Relative IQR (%) | Ratio vs PSNR (Discriminative Multiplier) |", + "|---|---|---|" +] + +for metric in METRICS: + if metric == "psnr_full": continue + d = cross_method_analysis["metrics"][metric] + rel_str = f"{d['relative_iqr_pct']:.2f}%" if d['relative_iqr_pct'] is not None else "N/A" + ratio_str = f"**{d['ratio_vs_psnr_relative']:.1f}x**" if d['ratio_vs_psnr_relative'] is not None else "N/A" + md_lines.append(f"| {metric} | {rel_str} | {ratio_str} |") + +md_content = "\n".join(md_lines) +with open(os.path.join(PROF_DIR, "part3b_findings.md"), "w") as f: + f.write(md_content) + +print(md_content) diff --git a/scripts/phase2_counterfactual/aggregate_part3b_supplement.py b/scripts/phase2_counterfactual/aggregate_part3b_supplement.py new file mode 100644 index 0000000000000000000000000000000000000000..e16bec339d2d1e0a563d5fbd4ece4a8f2cbe054a --- /dev/null +++ b/scripts/phase2_counterfactual/aggregate_part3b_supplement.py @@ -0,0 +1,112 @@ +import os, json, numpy as np + +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility" +PROF_DIR = os.path.join(OUTPUT_ROOT, "_profiles") + +METHODS = ['vanilla_3dgs', 'analyticsplatting', 'erankgs', 'ges', 'lightgaussian', 'minisplatting', 'opti3dgs', 'pgsr', 'steepgs'] +SCENES = ['bicycle','bonsai','counter','flowers','garden','kitchen','room','stump','treehill', + 'auditorium','ballroom','barn','caterpillar','courtroom','lighthouse','museum','palace','playground','temple','train','truck', + 'DrJohnson','Playroom', + 'Chair','Drums','Ficus','Hotdog','Lego','Materials','Mic','Ship'] +METRICS = ["psnr_full", "sh_net_effect", "sh_corruption_rate", "opacity_net_effect", "coverage", "coverage_error_fraction", "residual_error"] + +data = {m: {} for m in METHODS} +status_counts = {m: {"PASS": 0, "WARN": 0, "FAIL": 0} for m in METHODS} +failed_cells = [] + +# --- 1. Load Data & Status Statistics --- +for m in METHODS: + for s in SCENES: + p = os.path.join(OUTPUT_ROOT, f"{m}_{s}", "counterfactual_summary.json") + if os.path.exists(p): + with open(p) as f: + d = json.load(f) + data[m][s] = d + st = d.get("sanity_status", "FAIL") + status_counts[m][st] = status_counts[m].get(st, 0) + 1 + if st == "FAIL": + failed_cells.append(f"{m}_{s}: {d.get('sanity_failures', [])}") + +medians = {m: {} for m in METHODS} +for m in METHODS: + for k in METRICS: + vals = [data[m][s][k] for s in SCENES if s in data[m] and data[m][s].get(k) is not None] + medians[m][k] = float(np.median(vals)) if vals else 0.0 + +# --- 2. Saturated Subset Analysis --- +# 寻找 PSNR 最接近的 4 个方法构成的饱和簇 +sorted_by_psnr = sorted(METHODS, key=lambda x: medians[x]["psnr_full"]) +best_window = [] +min_psnr_range = 9999 + +for i in range(len(sorted_by_psnr) - 3): + window = sorted_by_psnr[i:i+4] + rng = medians[window[-1]]["psnr_full"] - medians[window[0]]["psnr_full"] + if rng < min_psnr_range: + min_psnr_range = rng + best_window = window + +def calc_iqr_stats(vals_list): + if not vals_list: return None, None, None + med = float(np.median(vals_list)) + abs_iqr = float(np.percentile(vals_list, 75) - np.percentile(vals_list, 25)) + rel_iqr = (abs_iqr / abs(med) * 100) if med != 0 else None + return med, abs_iqr, rel_iqr + +sat_psnr_vals = [medians[m]["psnr_full"] for m in best_window] +_, _, sat_psnr_rel_iqr = calc_iqr_stats(sat_psnr_vals) + +sat_results = {"subset_methods": best_window, "psnr_range_db": min_psnr_range, "psnr_rel_iqr": sat_psnr_rel_iqr, "metrics": {}} +for k in METRICS: + if k == "psnr_full": continue + vals = [medians[m][k] for m in best_window] + _, _, rel_iqr = calc_iqr_stats(vals) + sat_results["metrics"][k] = { + "rel_iqr": rel_iqr, + "ratio_vs_psnr": (rel_iqr / sat_psnr_rel_iqr) if (rel_iqr is not None and sat_psnr_rel_iqr) else None + } + +with open(os.path.join(PROF_DIR, "psnr_saturated_subset_analysis.json"), "w") as f: + json.dump(sat_results, f, indent=4) + +# --- 3. Stratification (Delta Signs vs Vanilla) --- +strat_results = {} +for m in METHODS: + if m == "vanilla_3dgs": continue + strat_results[m] = {} + for k in METRICS: + if k == "psnr_full": continue + delta = medians[m][k] - medians["vanilla_3dgs"][k] + strat_results[m][k] = {"delta": delta, "sign": "+" if delta > 0 else ("-" if delta < 0 else "0")} + +with open(os.path.join(PROF_DIR, "method_family_stratified.json"), "w") as f: + json.dump(strat_results, f, indent=4) + +# --- Output Report --- +print("=== 1. Execution Sanity Status ===") +print(f"{'Method':<20} | {'PASS':<6} | {'WARN':<6} | {'FAIL':<6}") +print("-" * 45) +for m in METHODS: + print(f"{m:<20} | {status_counts[m]['PASS']:<6} | {status_counts[m]['WARN']:<6} | {status_counts[m]['FAIL']:<6}") +print(f"\nTotal FAILS across all {len(METHODS)*len(SCENES)} cells: {len(failed_cells)}") + +print(f"\n=== 2. PSNR Saturated Subset Analysis ===") +print(f"Subset identified: {best_window}") +print(f"PSNR Relative IQR of this subset: {sat_psnr_rel_iqr:.2f}% (Range: {min_psnr_range:.2f} dB)") +print(f"\n{'Metric':<25} | {'Rel IQR (%)':<15} | {'Ratio vs PSNR':<15}") +print("-" * 60) +for k in METRICS: + if k == "psnr_full": continue + d = sat_results["metrics"][k] + rel = f"{d['rel_iqr']:.2f}%" if d['rel_iqr'] else "N/A" + rat = f"{d['ratio_vs_psnr']:.1f}x" if d['ratio_vs_psnr'] else "N/A" + print(f"{k:<25} | {rel:<15} | {rat:<15}") + +print(f"\n=== 3. Family Stratification (Sign vs Vanilla) ===") +print(f"{'Method':<20} | {'sh_net':<8} | {'sh_corr':<8} | {'opa_net':<8} | {'cov_err':<8} | {'resid':<8}") +print("-" * 70) +for m in METHODS: + if m == "vanilla_3dgs": continue + d = strat_results[m] + print(f"{m:<20} | {d['sh_net_effect']['sign']:<8} | {d['sh_corruption_rate']['sign']:<8} | {d['opacity_net_effect']['sign']:<8} | {d['coverage_error_fraction']['sign']:<8} | {d['residual_error']['sign']:<8}") + diff --git a/scripts/phase2_counterfactual/counterfactual_single.py b/scripts/phase2_counterfactual/counterfactual_single.py new file mode 100644 index 0000000000000000000000000000000000000000..f4a9cd0786a3bc3bac85dd065d1c0d3eda263cdb --- /dev/null +++ b/scripts/phase2_counterfactual/counterfactual_single.py @@ -0,0 +1,142 @@ +import os, sys, math, json, torch, numpy as np, gsplat + +sys.path.insert(0, "/root/autodl-tmp/SplatAtlas/scripts/phase1_validation") +sys.path.insert(0, "/root/autodl-tmp/3dgsAtlas_official") + +try: + from render_single import load_ply, load_test_cameras + from utils.general_utils import PILtoTorch +except ImportError as e: + print(f"[CRITICAL] ImportError: {e}") + sys.exit(1) + +def render(tuple_, cam, bg_color): + means, quats, scales, opacities, shs = tuple_ + device = means.device + viewmat = torch.tensor(cam['viewmat'], device=device).unsqueeze(0) + K = torch.tensor(cam['K'], device=device).unsqueeze(0) + colors, alphas, _ = gsplat.rasterization( + means=means, quats=quats, scales=scales, opacities=opacities, colors=shs, + viewmats=viewmat, Ks=K, width=cam['width'], height=cam['height'], + sh_degree=3, packed=True, render_mode='RGB', backgrounds=bg_color.unsqueeze(0), + ) + return colors.squeeze(0), alphas.squeeze(0) + +def load_gt(cam, bg_color_chw, device): + gt_chw_full = PILtoTorch(cam['pil_image'], (cam['width'], cam['height'])).to(device) + if gt_chw_full.shape[0] == 4: + alpha = gt_chw_full[3:4, ...] + rgb = gt_chw_full[:3, ...] + gt_chw = rgb * alpha + bg_color_chw * (1.0 - alpha) + else: + gt_chw = gt_chw_full[:3, ...] + return gt_chw.permute(1, 2, 0) + +def run_single_cell(ply_path: str, source_path: str, resolution: int, bg_rgb: list, sh_degree: int, out_json_path: str, phase1_ref_psnr: float = None, dataset: str = "unknown", method: str = "vanilla_3dgs", part_id: str = "2") -> dict: + device = 'cuda' + bg_color = torch.tensor(bg_rgb, dtype=torch.float32, device=device) + bg_color_chw = bg_color.view(3, 1, 1) + + means, quats, scales, opacities, shs = load_ply(ply_path) + cameras = load_test_cameras(source_path, resolution) + n_views = len(cameras) + if n_views == 0: raise ValueError("No test cameras found!") + + full_tuple = (means, quats, scales, opacities, shs) + sh_zero_shs = shs.clone(); sh_zero_shs[:, 1:, :] = 0.0 + sh_zero_tuple = (means, quats, scales, opacities, sh_zero_shs) + opa_one_opa = torch.ones_like(opacities) + opa_one_tuple = (means, quats, scales, opa_one_opa, shs) + + tau = 0.0 + all_E_full, all_E_DC, all_E_opa1, all_T_rem, all_C_bg_diff = [], [], [], [], [] + psnr_per_view, mse_accum_pooled, n_pixels_total = [], 0.0, 0 + + with torch.no_grad(): + for cam in cameras: + C_full, alpha_full = render(full_tuple, cam, bg_color) + C_full = C_full.clamp(0.0, 1.0) + C_DC, _ = render(sh_zero_tuple, cam, bg_color) + C_DC = C_DC.clamp(0.0, 1.0) + C_opa1, _ = render(opa_one_tuple, cam, bg_color) + C_opa1 = C_opa1.clamp(0.0, 1.0) + + gt_img = load_gt(cam, bg_color_chw, device) + H, W = cam['height'], cam['width'] + n_pixels_total += (H * W) + + mse_view = ((C_full - gt_img)**2).mean().item() + psnr_v = 10 * math.log10(1.0 / mse_view) if mse_view > 0 else 100.0 + psnr_per_view.append(psnr_v) + mse_accum_pooled += mse_view * (H * W) + + all_E_full.append((C_full - gt_img).abs().mean(dim=-1).flatten().cpu()) + all_E_DC.append((C_DC - gt_img).abs().mean(dim=-1).flatten().cpu()) + all_E_opa1.append((C_opa1 - gt_img).abs().mean(dim=-1).flatten().cpu()) + all_T_rem.append((1.0 - alpha_full.squeeze(-1)).flatten().cpu()) + all_C_bg_diff.append((bg_color.view(1, 1, 3) - gt_img).abs().mean(dim=-1).flatten().cpu()) + + E_full = torch.cat(all_E_full); E_DC = torch.cat(all_E_DC); E_opa1 = torch.cat(all_E_opa1) + T_rem = torch.cat(all_T_rem); C_bg_diff = torch.cat(all_C_bg_diff) + + psnr_full = float(np.mean(psnr_per_view)) + psnr_full_pooled = 10 * math.log10(1.0 / (mse_accum_pooled / n_pixels_total)) + + coverage = 1.0 - T_rem.mean().item() + coverage_error_fraction = ((T_rem * C_bg_diff).sum() / E_full.sum()).item() + dE_SH = E_DC - E_full; dE_opa = E_opa1 - E_full + + sh_net_effect = dE_SH.mean().item() + sh_corr_mask = dE_SH < -tau + sh_corruption_rate = sh_corr_mask.float().mean().item() + sh_corruption_severity = dE_SH[sh_corr_mask].mean().item() if sh_corr_mask.any() else None + + opacity_net_effect = dE_opa.mean().item() + opacity_pathology_rate = (dE_opa < -tau).float().mean().item() + + res_mask = (dE_SH > 0) & (dE_opa > 0) & (T_rem < 0.05) + residual_error = E_full[res_mask].mean().item() if res_mask.any() else 0.0 + + fails = [] + sh_negative_flag = bool(sh_net_effect < 0) + + if part_id in ["3a", "3b"]: + if phase1_ref_psnr is not None: + diff = abs(psnr_full - phase1_ref_psnr) + if diff > 2.0: fails.append("psnr_deviation_exceeds_2.0db") + elif diff > 0.8: fails.append("psnr_deviation_exceeds_0.8db") + if opacity_net_effect < -0.05: fails.append("opacity_net_effect_below_minus_0.05") + if coverage < 0.80 and dataset != "synthetic": fails.append("coverage_below_0.80") + + if len(fails) == 0: sanity_status = "PASS" + elif len(fails) == 1 and "psnr_deviation_exceeds_0.8db" in fails: sanity_status = "WARN" + else: sanity_status = "FAIL" + else: + if phase1_ref_psnr is not None and abs(psnr_full - phase1_ref_psnr) > 0.5: fails.append("psnr_deviation_exceeds_0.5db") + if sh_net_effect <= 0: fails.append("sh_net_effect_not_positive") + if opacity_net_effect < -0.01: fails.append("opacity_net_effect_below_minus_0.01") + if coverage < 0.85 and dataset != "synthetic": fails.append("coverage_below_0.85") + + if len(fails) == 0: sanity_status = "PASS" + elif len(fails) == 1 and fails[0] == "psnr_deviation_exceeds_0.5db" and phase1_ref_psnr is not None and abs(psnr_full - phase1_ref_psnr) <= 1.0: sanity_status = "WARN" + else: sanity_status = "FAIL" + + out_data = { + "method": method, "scene": os.path.basename(os.path.dirname(out_json_path)).replace(f"{method}_", ""), + "part": part_id, "sh_negative_flag": sh_negative_flag, + "dataset": dataset, "n_views": n_views, "n_pixels_total": n_pixels_total, + "psnr_full": psnr_full, "psnr_full_pooled": psnr_full_pooled, "phase1_reference_psnr": phase1_ref_psnr, + "coverage": coverage, "coverage_error_fraction": coverage_error_fraction, + "sh_net_effect": sh_net_effect, "sh_corruption_rate": sh_corruption_rate, "sh_corruption_severity": sh_corruption_severity, + "opacity_net_effect": opacity_net_effect, "opacity_pathology_rate": opacity_pathology_rate, + "residual_error": residual_error, "noise_floor_tau": tau, "background_color": bg_rgb, + "sanity_status": sanity_status, "sanity_failures": fails + } + if phase1_ref_psnr is None: out_data["sanity_check_1_note"] = "skipped: no reference" + + os.makedirs(os.path.dirname(out_json_path), exist_ok=True) + with open(out_json_path, 'w') as f: json.dump(out_data, f, indent=4) + return out_data + +if __name__ == "__main__": + pass diff --git a/scripts/phase2_counterfactual/counterfactual_vanilla_batch.py b/scripts/phase2_counterfactual/counterfactual_vanilla_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..fd1cec94aca2000af22ee86260cca60ed0ba02ee --- /dev/null +++ b/scripts/phase2_counterfactual/counterfactual_vanilla_batch.py @@ -0,0 +1,125 @@ +import os, sys, json, csv, argparse, time, traceback +from datetime import datetime, timezone + +sys.path.insert(0, "/root/autodl-tmp/SplatAtlas/scripts/phase2_counterfactual") +from counterfactual_single import run_single_cell + +SCENES = [ + "bicycle", "bonsai", "counter", "flowers", "garden", "kitchen", "room", "stump", "treehill", + "auditorium", "ballroom", "barn", "caterpillar", "courtroom", "lighthouse", "museum", "palace", "playground", "temple", "train", "truck", + "DrJohnson", "Playroom", + "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship" +] +MIP360_SET = {"bicycle", "bonsai", "counter", "flowers", "garden", "kitchen", "room", "stump", "treehill"} +TNT_SET = {"auditorium", "ballroom", "barn", "caterpillar", "courtroom", "lighthouse", "museum", "palace", "playground", "temple", "train", "truck"} +DB_SET = {"DrJohnson", "Playroom"} +SYNTHETIC_SET = {"Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"} + +DATASET_ROOTS = { + "mip360": "/root/autodl-tmp/dataset/360", "tnt": "/root/autodl-tmp/dataset/tnt", + "db": "/root/autodl-tmp/dataset/deepblending_clean", "synthetic": "/root/autodl-tmp/dataset/nerf_synthetic_off/nerf_synthetic" +} +BG_COLORS = { "mip360": [0.0, 0.0, 0.0], "tnt": [0.0, 0.0, 0.0], "db": [0.0, 0.0, 0.0], "synthetic": [1.0, 1.0, 1.0] } + +CSV_PATH = "/root/autodl-tmp/SplatAtlas/scripts/phase1_validation/validation_results_final.csv" +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility" +MANIFEST_PATH = os.path.join(OUTPUT_ROOT, "_manifest.json") +ERROR_LOG_PATH = os.path.join(OUTPUT_ROOT, "vanilla_batch_errors.log") + +def scene_to_dataset(name): + if name in MIP360_SET: return "mip360" + if name in TNT_SET: return "tnt" + if name in DB_SET: return "db" + if name in SYNTHETIC_SET: return "synthetic" + raise AssertionError(f"Unknown scene: {name}") + +def load_phase1_csv(): + csv_data = {} + if not os.path.exists(CSV_PATH): return csv_data + with open(CSV_PATH, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + m = row['method'].strip() + s = row['scene'].strip().lower() + res_val = int(row['resolution']) if row['resolution'] else None + psnr_val = float(row['gsplat_psnr']) if row['gsplat_psnr'] else None + csv_data[(m, s)] = { "resolution": res_val, "gsplat_psnr": psnr_val, "status": row['status'] } + return csv_data + +def run_dry_run(target_methods, target_scenes): + csv_data = load_phase1_csv() + manifest_cells = [] + + print(f"{'Method':<15} | {'Scene':<15} | {'Dataset':<10} | {'PLY_OK':<6} | {'SRC_OK':<6}") + print("-" * 65) + + for method in target_methods: + for scene in target_scenes: + ds = scene_to_dataset(scene) + scene_lower = scene.lower() + + csv_info = csv_data.get((method, scene_lower), None) + phase1_in_csv = csv_info is not None + + if phase1_in_csv: + res, ref_psnr, phase1_status = csv_info["resolution"], csv_info["gsplat_psnr"], csv_info["status"] + else: + res, ref_psnr, phase1_status = 1, None, "UNKNOWN" + + bg = BG_COLORS[ds] + ply_path = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}/point_cloud/iteration_30000/point_cloud.ply" + src_path = os.path.join(DATASET_ROOTS[ds], scene.lower() if ds == "synthetic" else scene) + out_json = os.path.join(OUTPUT_ROOT, f"{method}_{scene}", "counterfactual_summary.json") + + ply_ok, src_ok = os.path.exists(ply_path), os.path.exists(src_path) + + manifest_cells.append({ + "method": method, "scene": scene, "dataset": ds, "resolution": res, "bg_rgb": bg, + "ply_path": ply_path, "ply_ok": ply_ok, "source_path": src_path, "src_ok": src_ok, + "out_json_path": out_json, "phase1_ref_psnr": ref_psnr, "phase1_status": phase1_status, "phase1_in_csv": phase1_in_csv + }) + print(f"{method:<15} | {scene:<15} | {ds:<10} | {str(ply_ok):<6} | {str(src_ok):<6}") + + print("-" * 65) + print(f"[STAT] Total cells: {len(manifest_cells)}") + + os.makedirs(OUTPUT_ROOT, exist_ok=True) + with open(MANIFEST_PATH, 'w') as f: + json.dump({"generated_at": datetime.now(timezone.utc).isoformat(), "n_cells": len(manifest_cells), "cells": manifest_cells}, f, indent=4) + print(f"\n[Dry Run Complete] Manifest written to {MANIFEST_PATH}") + +def run_execute(): + if not os.path.exists(MANIFEST_PATH): sys.exit("[ERROR] Manifest not found.") + with open(MANIFEST_PATH, 'r') as f: manifest = json.load(f) + + start_time = time.time() + cells = manifest["cells"] + + for i, cell in enumerate(cells): + method, scene = cell["method"], cell["scene"] + part_id = "2" if method == "vanilla_3dgs" else ("3a" if method in ["3dgsmcmc", "absgs", "gaussianpro", "pixelgs"] else "3b") + t_elapsed = time.time() - start_time + eta = (len(cells) - i) * (t_elapsed / i) if i > 0 else 0 + + print(f"[{i+1}/{len(cells)}] {method} · {scene} | elapsed={t_elapsed:.1f}s | eta={eta:.1f}s") + try: + run_single_cell( + ply_path=cell["ply_path"], source_path=cell["source_path"], resolution=cell["resolution"], + bg_rgb=cell["bg_rgb"], sh_degree=3, out_json_path=cell["out_json_path"], + phase1_ref_psnr=cell["phase1_ref_psnr"], dataset=cell["dataset"], method=method, part_id=part_id + ) + except Exception as e: + print(f" -> [FAIL] Exception isolated. Check errors.log") + with open(ERROR_LOG_PATH, "a") as f: f.write(f"[{datetime.now().isoformat()}] Method: {method}, Scene: {scene}\n{traceback.format_exc()}\n") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--dry-run", action="store_true") + group.add_argument("--execute", action="store_true") + parser.add_argument("--methods", type=str, default="vanilla_3dgs") + parser.add_argument("--scenes", type=str, default="ALL") + args = parser.parse_args() + + if args.dry_run: run_dry_run(args.methods.split(','), SCENES if args.scenes == "ALL" else args.scenes.split(',')) + elif args.execute: run_execute() diff --git a/scripts/phase2_counterfactual/fix_bonsai_res.py b/scripts/phase2_counterfactual/fix_bonsai_res.py new file mode 100644 index 0000000000000000000000000000000000000000..c4b8d148f3f73604282d56379f785021e39e1339 --- /dev/null +++ b/scripts/phase2_counterfactual/fix_bonsai_res.py @@ -0,0 +1,24 @@ +import os +from PIL import Image +from concurrent.futures import ThreadPoolExecutor + +scene_dir = "/root/autodl-tmp/dataset/360/bonsai" +img_dir = os.path.join(scene_dir, "images") +target_dir = os.path.join(scene_dir, "images_2") + +if not os.path.exists(target_dir): + os.makedirs(target_dir) + images = [f for f in os.listdir(img_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] + + def process(img_name): + src = os.path.join(img_dir, img_name) + dst = os.path.join(target_dir, img_name) + with Image.open(src) as img: + img.resize((img.width // 2, img.height // 2), Image.Resampling.LANCZOS).save(dst) + + print(f"[Data Alignment] Downsampling {len(images)} images for bonsai (resolution=2)...") + with ThreadPoolExecutor() as pool: + pool.map(process, images) + print("[Data Alignment] bonsai images_2 generation complete.") +else: + print("[Data Alignment] bonsai images_2 already exists.") diff --git a/scripts/phase2_counterfactual/fix_resolutions.py b/scripts/phase2_counterfactual/fix_resolutions.py new file mode 100644 index 0000000000000000000000000000000000000000..186428a502e67d6810d83f20d4f448fbdbae60d2 --- /dev/null +++ b/scripts/phase2_counterfactual/fix_resolutions.py @@ -0,0 +1,61 @@ +import os +from PIL import Image +from concurrent.futures import ThreadPoolExecutor +from tqdm import tqdm + +DATASETS = [ + "/root/autodl-tmp/dataset/360", + "/root/autodl-tmp/dataset/tnt" +] + +def process_image(src_path, dst_path, scale): + if os.path.exists(dst_path): + return + try: + with Image.open(src_path) as img: + new_size = (int(img.width / scale), int(img.height / scale)) + resized_img = img.resize(new_size, Image.Resampling.LANCZOS) + resized_img.save(dst_path) + except Exception as e: + print(f"Error processing {src_path}: {e}") + +def downsample_scene(scene_dir): + img_dir = os.path.join(scene_dir, "images") + if not os.path.exists(img_dir): + return + + images = [f for f in os.listdir(img_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] + + # 需要生成的降采样层级 + scales = [2, 4] + + tasks = [] + for scale in scales: + target_dir = os.path.join(scene_dir, f"images_{scale}") + os.makedirs(target_dir, exist_ok=True) + + for img_name in images: + src_path = os.path.join(img_dir, img_name) + dst_path = os.path.join(target_dir, img_name) + tasks.append((src_path, dst_path, scale)) + + return tasks + +if __name__ == "__main__": + all_tasks = [] + print("[Preprocessing] Scanning datasets for missing resolution folders...") + for ds_root in DATASETS: + if not os.path.exists(ds_root): continue + for scene in os.listdir(ds_root): + scene_dir = os.path.join(ds_root, scene) + if os.path.isdir(scene_dir): + tasks = downsample_scene(scene_dir) + if tasks: all_tasks.extend(tasks) + + if not all_tasks: + print("[Preprocessing] All resolution folders exist. Nothing to do.") + else: + print(f"[Preprocessing] Generating {len(all_tasks)} downsampled images. This will use all CPU cores...") + with ThreadPoolExecutor() as executor: + list(tqdm(executor.map(lambda p: process_image(*p), all_tasks), total=len(all_tasks))) + print("[Preprocessing] Done!") diff --git a/scripts/phase2_counterfactual/part3a_reanalysis_generator.py b/scripts/phase2_counterfactual/part3a_reanalysis_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..5e3546a59f39b6cd7fbc8e7c198b8c0d3cb8c381 --- /dev/null +++ b/scripts/phase2_counterfactual/part3a_reanalysis_generator.py @@ -0,0 +1,184 @@ +import os, json, numpy as np +from datetime import datetime, timezone + +OUTPUT_ROOT = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility" +PROF_DIR = os.path.join(OUTPUT_ROOT, "_profiles") +METHODS = ["vanilla_3dgs", "3dgsmcmc", "absgs", "gaussianpro", "pixelgs"] +INCLUDED_METHODS = ["vanilla_3dgs", "absgs", "gaussianpro", "pixelgs"] +SCENES = ["bicycle", "bonsai", "counter", "flowers", "garden", "kitchen", "room", "stump", "treehill", + "auditorium", "ballroom", "barn", "caterpillar", "courtroom", "lighthouse", "museum", "palace", "playground", "temple", "train", "truck", + "DrJohnson", "Playroom", + "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"] +METRICS = ["psnr_full", "sh_net_effect", "sh_corruption_rate", "opacity_net_effect", "coverage", "coverage_error_fraction", "residual_error"] + +data = {m: {} for m in METHODS} +sh_neg_scenes = {"3dgsmcmc": [], "absgs": [], "gaussianpro": [], "pixelgs": []} + +for m in METHODS: + for s in SCENES: + p = os.path.join(OUTPUT_ROOT, f"{m}_{s}", "counterfactual_summary.json") + if os.path.exists(p): + with open(p) as f: + d = json.load(f) + data[m][s] = d + if m != "vanilla_3dgs" and d.get("sh_negative_flag"): + sh_neg_scenes[m].append(s) + +# Cross scene medians +per_method_cross_scene_median = {} +for m in METHODS: + per_method_cross_scene_median[m] = {} + for metric in METRICS: + vals = [data[m][s].get(metric) for s in SCENES if s in data[m] and data[m][s].get(metric) is not None] + per_method_cross_scene_median[m][metric] = float(np.median(vals)) if vals else 0.0 + +def calc_iqr_stats(val_dict): + vals = list(val_dict.values()) + if not vals: return None, None, None + median = float(np.median(vals)) + abs_iqr = float(np.percentile(vals, 75) - np.percentile(vals, 25)) + rel_iqr = (abs_iqr / abs(median) * 100) if median != 0 else None + return median, abs_iqr, rel_iqr + +# 1. Bonsai Analysis (5 Methods) +bonsai_analysis = {"n_methods": 5, "metrics": {}} +b_psnr_vals = {m: data[m]["bonsai"].get("psnr_full", 0) for m in METHODS if "bonsai" in data[m]} +_, psnr_abs_iqr_bonsai, psnr_rel_iqr_bonsai = calc_iqr_stats(b_psnr_vals) + +for metric in METRICS: + val_dict = {m: data[m]["bonsai"].get(metric, 0) for m in METHODS if "bonsai" in data[m]} + med, abs_iqr, rel_iqr = calc_iqr_stats(val_dict) + + ratio_abs = (abs_iqr / psnr_abs_iqr_bonsai) if psnr_abs_iqr_bonsai else None + ratio_rel = (rel_iqr / psnr_rel_iqr_bonsai) if (rel_iqr is not None and psnr_rel_iqr_bonsai) else None + + bonsai_analysis["metrics"][metric] = { + "values_5_methods": val_dict, + "absolute_iqr": abs_iqr, + "median_5_methods": med, + "relative_iqr_pct": rel_iqr, + "ratio_vs_psnr_absolute": ratio_abs, + "ratio_vs_psnr_relative": ratio_rel + } + +# 2. Cross Scene Analysis (4 included methods) +cross_scene_analysis = {"n_methods": 4, "methods_included": INCLUDED_METHODS, "metrics": {}} +cs_psnr_vals = {m: per_method_cross_scene_median[m]["psnr_full"] for m in INCLUDED_METHODS} +_, psnr_abs_iqr_cs, psnr_rel_iqr_cs = calc_iqr_stats(cs_psnr_vals) + +for metric in METRICS: + val_dict = {m: per_method_cross_scene_median[m][metric] for m in INCLUDED_METHODS} + med, abs_iqr, rel_iqr = calc_iqr_stats(val_dict) + + ratio_abs = (abs_iqr / psnr_abs_iqr_cs) if psnr_abs_iqr_cs else None + ratio_rel = (rel_iqr / psnr_rel_iqr_cs) if (rel_iqr is not None and psnr_rel_iqr_cs) else None + + cross_scene_analysis["metrics"][metric] = { + "values_4_methods": val_dict, + "absolute_iqr": abs_iqr, + "median_4_methods": med, + "relative_iqr_pct": rel_iqr, + "ratio_vs_psnr_absolute": ratio_abs, + "ratio_vs_psnr_relative": ratio_rel + } + +# Build JSON +json_out = { + "generated_at": datetime.now(timezone.utc).isoformat(), + "methods_analyzed": METHODS, + "methods_excluded_from_thesis": ["3dgsmcmc"], + "methods_excluded_reason": { + "3dgsmcmc": "partial-failure: sh_sign mixed on 3 scenes, pipeline known to crash on subset of scenes" + }, + "sh_negative_scenes_by_method": sh_neg_scenes, + "per_method_cross_scene_median": per_method_cross_scene_median, + "bonsai_single_cell_iqr_analysis": bonsai_analysis, + "cross_scene_median_iqr_analysis": cross_scene_analysis, + "thesis_status": { + "claim": "photometric loss is strong output constraint but weak representation constraint", + "part3a_verdict": "inconclusive — method sampling bias (same family), PSNR not saturated, absolute IQR scale-mixed", + "pending_validation": "Part 3b 9-method heterogeneous set with relative IQR and PSNR-saturation filtering" + } +} + +# Build Markdown +mcmc_neg_scenes = sh_neg_scenes.get("3dgsmcmc", []) +md_content = f"""# Phase 2 Part 3a — Findings and Caveats + +## 1. Scope +- Methods: vanilla_3dgs (reference), 3dgsmcmc*, absgs, gaussianpro, pixelgs +- Scenes: 31 (Mip360×9 + T&T×12 + DB×2 + Synthetic×8) +- Total cells: 124 new + 31 vanilla reference = 155 +- *3dgsmcmc flagged as partial-failure method; included in file-level outputs but excluded from thesis-level statistical analysis. + +## 2. Primary Findings + +### Finding 2.1: Convergent densification-level improvement +All four densification variants demonstrate a consistent, unidirectional improvement over the Vanilla baseline in terms of geometric coverage and reduced reliance on semi-transparent floaters. Specifically, geometric coverage improved consistently (+0.01 to +0.03 delta medians), opacity net effect decreased (-0.03 to -0.08 delta medians), and the coverage error fraction was reduced remarkably (-0.03 to -0.11 delta medians). This unidirectional movement across all variants is a strong, non-accidental signal of structural improvement over Vanilla 3DGS. + +### Finding 2.2: Scene-dependent opacity channel stratification (carried from Part 2) +The dataset-level stratification observed in Part 2 (e.g., Synthetic opacity net effect ~0.014 vs. T&T ~0.236, a 17x gap) is strictly preserved across the Part 3a variants. While the absolute opacity net effect is universally reduced by these methods, the structural hierarchy between datasets remains intact, proving that dataset characteristics govern the opacity error profile far more than the specific densification method used. + +### Finding 2.3: Weak counterfactual-channel discrimination among same-family methods +When evaluated strictly among these four same-family densification variants, the absolute IQR of the counterfactual channels remains tightly bounded. This confirms that methods sharing the underlying representational semantics (only altering densification timing/location) yield highly similar internal representations. This must be caveated as an observation "among same-family methods" rather than a generalized rule for all variants. + +## 3. Reanalysis of IQR Ratios + +This section re-evaluates the discriminative power using Relative IQR (`IQR / |median|`) to correct for the scale mismatch between dB (PSNR) and [0,1] domain metrics. + +**Bonsai Single-Cell (5 Methods)** +- PSNR Relative IQR: ~{bonsai_analysis['metrics']['psnr_full']['relative_iqr_pct']:.2f}% +- Coverage Error Fraction Relative IQR: ~{bonsai_analysis['metrics']['coverage_error_fraction']['relative_iqr_pct']:.2f}% +- *Result*: In relative terms, the dispersion of Coverage Error Fraction is significantly larger than PSNR, increasing the counterfactual channel's discriminative ratio substantially over the obsolete absolute scale. + +**Cross-Scene Median (4 Thesis-Included Methods)** +- PSNR Relative IQR: ~{cross_scene_analysis['metrics']['psnr_full']['relative_iqr_pct']:.2f}% +- Coverage Error Fraction Relative IQR: ~{cross_scene_analysis['metrics']['coverage_error_fraction']['relative_iqr_pct']:.2f}% + +## 4. Caveats + +### Caveat 4.1: Method sampling bias +The four methods evaluated in Part 3a are exclusively densification strategy variants. They do not alter the fundamental PLY, SH, or opacity semantics of Vanilla 3DGS. Mechanistically, their counterfactual profiles are expected to cluster tightly. The observation of similar representations among this cohort does not overturn the core thesis; rather, it validates the pipeline's logic. True thesis validation requires the heterogeneous 9-method set in Part 3b (incorporating quantization, rank-correlation, etc.). + +### Caveat 4.2: PSNR range is not saturated +The PSNR range across the 5 methods on pilot scenes (e.g., bonsai spanning 31.97 to 32.67) represents a 0.70 dB gap, translating to a ~1.17x difference in MSE. This does not meet the literature definition of strict PSNR saturation (typically IQR < 0.15 dB). The extraordinary discriminative power of counterfactual metrics must be tested in true saturation regimes, which Part 3a has not yet entered. + +### Caveat 4.3: Absolute IQR comparison is scale-mixed +The initial protocol comparing absolute IQR between a logarithmic scale (PSNR) and linear fractions [0,1] is inherently biased toward PSNR. Correct formulation requires utilizing Relative IQR or converting PSNR to linear MSE equivalents. All subsequent thesis evaluations will adopt the Relative IQR standard. + +### Caveat 4.4: 3dgsmcmc partial failure +Both `3dgsmcmc` and `reactgs` exhibit known corner-case failures in training/rendering pipelines on specific scene subsets. Specifically, `3dgsmcmc` produced inverted SH behaviors (sh_net_effect < 0) on {len(mcmc_neg_scenes)} specific scenes: {', '.join(mcmc_neg_scenes) if mcmc_neg_scenes else 'None'}. While its 124-cell outputs are archived, it is explicitly excluded from thesis-level statistical claims. `reactgs` will be removed from the Part 3b roster. + +## 5. Thesis Status +The core SplatAtlas thesis—that "photometric loss acts as a strong output constraint but a weak representation constraint"—is currently **inconclusive but not refuted**. Part 3a demonstrated that within a tightly coupled family of methods, representations *do* converge similarly to output metrics. + +However, Part 3a also yielded two robust, independent findings: the dataset-driven stratification of opacity and the method-level convergent improvement of geometric coverage. These prove the counterfactual pipeline successfully extracts sub-PSNR structural behaviors. The ultimate validation of the thesis is deferred to Part 3b, which will deploy a heterogeneous method array and utilize corrected Relative IQR scaling in strict PSNR-saturation subsets. + +## 6. Deliverables +- 124 per-cell summaries under `/root/autodl-tmp/SplatAtlas/outputs/error_responsibility/{{method}}_{{scene}}/` +- 4 method profiles under `/root/autodl-tmp/SplatAtlas/outputs/error_responsibility/_profiles/{{method}}_profile.json` +- This findings doc (`part3a_findings_and_caveats.md`) +- Reanalysis data (`part3a_reanalysis.json`) +""" + +with open(os.path.join(PROF_DIR, "part3a_findings_and_caveats.md"), "w") as f: + f.write(md_content) +with open(os.path.join(PROF_DIR, "part3a_reanalysis.json"), "w") as f: + json.dump(json_out, f, indent=4) + +print("=== 1. Markdown Content ===") +print(md_content) + +print("\n=== 2. JSON Content (Preview) ===") +print(json.dumps(json_out, indent=2)) + +print("\n=== 3. Old vs New Ratio Comparison (Bonsai Single Cell) ===") +print(f"{'Metric':<25} | {'Old Abs Ratio':<15} | {'New Rel Ratio':<15}") +print("-" * 60) +for metric in METRICS: + if metric == "psnr_full": continue + old = bonsai_analysis["metrics"][metric]["ratio_vs_psnr_absolute"] + new = bonsai_analysis["metrics"][metric]["ratio_vs_psnr_relative"] + old_str = f"{old:.4f}" if old is not None else "N/A" + new_str = f"{new:.4f}" if new is not None else "N/A" + print(f"{metric:<25} | {old_str:<15} | {new_str:<15}") diff --git a/scripts/phase2_counterfactual/patch_router.py b/scripts/phase2_counterfactual/patch_router.py new file mode 100644 index 0000000000000000000000000000000000000000..4a44b40d618eb922210549d1dc9e3314bd68e0b6 --- /dev/null +++ b/scripts/phase2_counterfactual/patch_router.py @@ -0,0 +1,12 @@ +import os +path = "/root/autodl-tmp/SplatAtlas/scripts/phase1_validation/render_single.py" +with open(path, "r") as f: code = f.read() +old_str = 'is_synthetic = os.path.exists(os.path.join(source_path, "transforms_test.json"))' +new_str = 'is_synthetic = os.path.exists(os.path.join(source_path, "transforms_test.json")) and not os.path.exists(os.path.join(source_path, "sparse"))' + +if old_str in code: + code = code.replace(old_str, new_str) + with open(path, "w") as f: f.write(code) + print("[Patch Applied] render_single.py routing logic hardened (sparse check added).") +else: + print("[Patch Skipped] old routing string not found.") diff --git a/scripts/phase2_counterfactual/print_table.py b/scripts/phase2_counterfactual/print_table.py new file mode 100644 index 0000000000000000000000000000000000000000..97f5faea6553638c080b7aae5a97f54644862fca --- /dev/null +++ b/scripts/phase2_counterfactual/print_table.py @@ -0,0 +1,32 @@ +import os, json + +methods = ["vanilla_3dgs", "3dgsmcmc", "absgs", "gaussianpro", "pixelgs"] + +header = "{:<15} | {:<9} | {:<9} | {:<9} | {:<9} | {:<8} | {:<8} | {:<9} | {:<6} | {}".format( + "method", "psnr_full", "sh_net", "sh_cor_rt", "opa_net", "coverage", "cov_err", "resid_err", "sanity", "sh_neg" +) +print(header) +print("-" * len(header)) + +for m in methods: + p = f"/root/autodl-tmp/SplatAtlas/outputs/error_responsibility/{m}_bonsai/counterfactual_summary.json" + if not os.path.exists(p): + print(f"{m:<15} | MISSING") + continue + + with open(p) as f: + d = json.load(f) + + row = "{:<15} | {:<9.4f} | {:<9.4f} | {:<9.4f} | {:<9.4f} | {:<8.4f} | {:<8.4f} | {:<9.4f} | {:<6} | {}".format( + m, + d.get("psnr_full", 0), + d.get("sh_net_effect", 0), + d.get("sh_corruption_rate", 0), + d.get("opacity_net_effect", 0), + d.get("coverage", 0), + d.get("coverage_error_fraction", 0), + d.get("residual_error", 0), + d.get("sanity_status", "N/A"), + d.get("sh_negative_flag", "N/A") + ) + print(row) diff --git a/scripts/phase2_counterfactual/run_recon.sh b/scripts/phase2_counterfactual/run_recon.sh new file mode 100644 index 0000000000000000000000000000000000000000..b5d9f38c0ddce9d4c3f85ef05a5cfd346b38fc2f --- /dev/null +++ b/scripts/phase2_counterfactual/run_recon.sh @@ -0,0 +1,59 @@ +#!/bin/bash +OUT="/root/autodl-tmp/SplatAtlas/outputs/phase1/phase1_recon.md" +mkdir -p /root/autodl-tmp/SplatAtlas/outputs/phase1 +> $OUT + +echo "# 步骤 2：扫描 mainline 输出目录" >> $OUT +cd /root/autodl-tmp/SplatAtlas/outputs +for d in mainline_2_6 mainline_3_5 mainline_4_7 mainline_9_10 mainline_D12_8; do + echo "=== $d ===" >> $OUT + if [ -d "$d" ]; then + find "$d" -maxdepth 3 \( -name "*.json" -o -name "*.csv" -o -name "*.png" \) -exec ls -lh {} \; >> $OUT + else + echo "MISSING: 目录不存在 ($d)" >> $OUT + fi +done + +echo -e "\n# 步骤 3：逐一核查 Task 1.1 目标 figure" >> $OUT +for f in "mainline_4_7/sh_geo_correlation_scatter.png" "mainline_3_5/ccf_lag_histograms.png" "mainline_3_5/ccf_by_method_heatmap.png" "mainline_D12_8/shape_umap_by_cluster.png" "mainline_D12_8/shape_umap_by_category.png" "mainline_D12_8/shape_dendrogram.png" "mainline_2_6/mantel_scatter.png" "mainline_2_6/mantel_by_dimension.png" "mainline_2_6/mantel_by_scene.png"; do + if [ -f "$f" ]; then + ls -lh "$f" >> $OUT + else + echo "MISSING: $f" >> $OUT + fi +done + +echo -e "\n# 步骤 4：逐一核查 Task 1.2 目标数据文件" >> $OUT +for f in "mainline_3_5/causal_order_summary.json" "mainline_3_5/ccf_results.csv" "mainline_D12_8/shape_clustering_metrics.json" "mainline_D12_8/reclassified_methods.json" "mainline_2_6/mantel_results.json" "mainline_3_5/shock_significance.json" "mainline_4_7/population_turnover.json" "mainline_9_10/adaptability_results.json"; do + echo "=== $f ===" >> $OUT + if [ -f "$f" ]; then + ls -lh "$f" >> $OUT + head -c 1500 "$f" >> $OUT + echo -e "\n---" >> $OUT + else + echo "MISSING: $f" >> $OUT + fi +done + +echo -e "\n# 步骤 5：核查 Task 1.3b 数据源" >> $OUT +echo "=== mainline_4_7/sh_buffer_analysis.csv 列名 ===" >> $OUT +[ -f "mainline_4_7/sh_buffer_analysis.csv" ] && head -2 mainline_4_7/sh_buffer_analysis.csv >> $OUT || echo "MISSING: mainline_4_7/sh_buffer_analysis.csv" >> $OUT +echo "=== mainline_4_7/sh_buffer_summary.json ===" >> $OUT +[ -f "mainline_4_7/sh_buffer_summary.json" ] && cat mainline_4_7/sh_buffer_summary.json >> $OUT || echo "MISSING: mainline_4_7/sh_buffer_summary.json" >> $OUT +DATA_DIR="/root/autodl-tmp/SplatAtlas/analysis/data" +echo "=== level_b_scalars.parquet ===" >> $OUT +[ -f "$DATA_DIR/level_b_scalars.parquet" ] && ls -lh "$DATA_DIR/level_b_scalars.parquet" >> $OUT || echo "MISSING: $DATA_DIR/level_b_scalars.parquet" >> $OUT +echo "=== step1_metrics_no_collapse.csv 列名 ===" >> $OUT +[ -f "$DATA_DIR/step1_metrics_no_collapse.csv" ] && head -2 "$DATA_DIR/step1_metrics_no_collapse.csv" >> $OUT || echo "MISSING: $DATA_DIR/step1_metrics_no_collapse.csv" >> $OUT + +echo -e "\n# 步骤 6：核查 config/methods.yaml" >> $OUT +wc -l /root/autodl-tmp/SplatAtlas/config/methods.yaml >> $OUT 2>/dev/null || echo "MISSING: /root/autodl-tmp/SplatAtlas/config/methods.yaml" >> $OUT +grep -c "name:" /root/autodl-tmp/SplatAtlas/config/methods.yaml >> $OUT 2>/dev/null || echo "MISSING: name_field_in_methods.yaml" >> $OUT +grep "year:" /root/autodl-tmp/SplatAtlas/config/methods.yaml | head -5 >> $OUT 2>/dev/null || echo "MISSING: year_field_in_methods.yaml" >> $OUT + +echo -e "\nRECON COMPLETE" >> $OUT + +# 附加 MISSING FILES 汇总段落 +echo -e "\n## MISSING FILES" >> $OUT +grep -oP "(?<=MISSING: ).*" $OUT >> $OUT +if ! grep -q "MISSING:" $OUT; then echo "None" >> $OUT; fi diff --git a/scripts/phase2_counterfactual/stage1_diagnostic.py b/scripts/phase2_counterfactual/stage1_diagnostic.py new file mode 100644 index 0000000000000000000000000000000000000000..0b13da68946665ec5dd356bfd81947ba2921373e --- /dev/null +++ b/scripts/phase2_counterfactual/stage1_diagnostic.py @@ -0,0 +1,125 @@ +import os, sys, json, glob +import numpy as np +from datetime import datetime, timezone + +def main(): + summary_paths = glob.glob("/root/autodl-tmp/SplatAtlas/outputs/error_responsibility/vanilla_3dgs_*/counterfactual_summary.json") + + if len(summary_paths) != 31: + print(f"[CRITICAL] Expected 31 summaries, found {len(summary_paths)}.") + os.system("ls -la /root/autodl-tmp/SplatAtlas/outputs/error_responsibility/vanilla_3dgs_*/") + sys.exit(1) + + data_records = [] + for p in summary_paths: + with open(p, 'r') as f: + d = json.load(f) + if "opacity_net_effect" not in d or "sh_net_effect" not in d: + print(f"[CRITICAL] Missing fields in {p}") + print(json.dumps(d, indent=2)) + sys.exit(1) + data_records.append({ + "scene": d.get("scene", "unknown"), + "dataset": d.get("dataset", "unknown"), + "opacity_net_effect": d["opacity_net_effect"], + "sh_net_effect": d["sh_net_effect"] + }) + + def calc_stats(key): + vals = [r[key] for r in data_records] + sorted_records = sorted([{"scene": r["scene"], "dataset": r["dataset"], "value": r[key]} for r in data_records], key=lambda x: x["value"]) + p25 = float(np.percentile(vals, 25)) + p75 = float(np.percentile(vals, 75)) + + per_dataset = {} + for ds in ["mip360", "tnt", "db", "synthetic"]: + ds_vals = [r[key] for r in data_records if r["dataset"] == ds] + per_dataset[ds] = float(np.median(ds_vals)) if ds_vals else None + + return { + "min": float(np.min(vals)), + "max": float(np.max(vals)), + "median": float(np.median(vals)), + "p25": p25, + "p75": p75, + "iqr": p75 - p25, + "std": float(np.std(vals)), + "per_dataset_median": per_dataset, + "per_scene_sorted": sorted_records + } + + op_stats = calc_stats("opacity_net_effect") + sh_stats = calc_stats("sh_net_effect") + + op_iqr = op_stats["iqr"] + sh_iqr = sh_stats["iqr"] + + if op_iqr > 0.02 and sh_iqr > 0.02: + cat = "property" + act = "proceed_to_stage_2" + note = "Healthy cross-scene variance for both Opacity and SH." + elif 0.005 <= op_iqr <= 0.02 and sh_iqr > 0.02: + cat = "near_constant_property" + act = "proceed_with_caveat" + note = "Opacity variance is tight but normal; SH variance is healthy." + elif op_iqr < 0.005 and sh_iqr > 0.02: + cat = "bug_suspected" + act = "halt_debug" + note = "Opacity IQR abnormally tight. Suspected opacity mutation truncation/masking." + elif op_iqr > 0.02 and sh_iqr < 0.005: + cat = "sh_bug_suspected" + act = "halt_debug" + note = "SH IQR abnormally tight. Theoretical anomaly requiring explicit debug." + elif op_iqr < 0.005 and sh_iqr < 0.005: + cat = "bug_suspected" + act = "halt_debug" + note = "Dual collapse in variance. Aggregation-level bug suspected." + else: + cat = "near_constant_property" + act = "proceed_with_caveat" + note = "Boundary condition hit. Proceeding with caution." + + verdict = { + "opacity_iqr": op_iqr, + "sh_iqr": sh_iqr, + "category": cat, + "recommended_action": act, + "notes": note + } + + out_dict = { + "generated_at": datetime.now(timezone.utc).isoformat(), + "n_scenes": 31, + "opacity_net_effect": op_stats, + "sh_net_effect": sh_stats, + "verdict": verdict + } + + out_dir = "/root/autodl-tmp/SplatAtlas/outputs/error_responsibility/_diagnostics" + os.makedirs(out_dir, exist_ok=True) + out_path = os.path.join(out_dir, "opacity_stability_check.json") + + with open(out_path, 'w') as f: + json.dump(out_dict, f, indent=4) + + print("=== Opacity Stability Diagnostic ===") + print(f"n_scenes: {out_dict['n_scenes']}") + print(f"opacity_net_effect: min={op_stats['min']:.6f} max={op_stats['max']:.6f} median={op_stats['median']:.6f} iqr={op_iqr:.6f} std={op_stats['std']:.6f}") + print(f"sh_net_effect: min={sh_stats['min']:.6f} max={sh_stats['max']:.6f} median={sh_stats['median']:.6f} iqr={sh_iqr:.6f} std={sh_stats['std']:.6f}\n") + + print("per_dataset opacity median:") + for ds in ["mip360", "tnt", "db", "synthetic"]: + v = op_stats['per_dataset_median'][ds] + print(f" {ds:<9}: {f'{v:.6f}' if v is not None else 'None'}") + + print("\nper_dataset sh median:") + for ds in ["mip360", "tnt", "db", "synthetic"]: + v = sh_stats['per_dataset_median'][ds] + print(f" {ds:<9}: {f'{v:.6f}' if v is not None else 'None'}") + + print(f"\nverdict: {cat}") + print(f"action : {act}") + print(f"note : {note}") + +if __name__ == "__main__": + main() diff --git a/scripts/plot_c1_c2_figures.py b/scripts/plot_c1_c2_figures.py new file mode 100644 index 0000000000000000000000000000000000000000..15ede3bf02bad89ff0525661b7bea962c09b003a --- /dev/null +++ b/scripts/plot_c1_c2_figures.py @@ -0,0 +1,58 @@ +import pandas as pd +import matplotlib.pyplot as plt +import seaborn as sns +from pathlib import Path + +# 设置学术图表样式 +plt.rcParams.update({'font.size': 12, 'pdf.fonttype': 42, 'ps.fonttype': 42}) +sns.set_theme(style="whitegrid", context="paper") + +OUT_DIR = Path("/root/autodl-tmp/SplatAtlas/outputs/surfel_hypothesis_7tnt_all/") +df_method = pd.read_csv(OUT_DIR / "method_mean_7tnt.csv") +df_scene = pd.read_csv(OUT_DIR / "scene_mean_7tnt.csv") + +# 过滤掉 V3 为空的行（比如 2DGS） +df_v3 = df_method.dropna(subset=['v3_normal_median_deg', 'center_pca_normal_median_deg']).copy() + +# ========================================== +# Figure 1: V3 vs Center-PCA Scatter Plot +# ========================================== +plt.figure(figsize=(6, 5)) +ax = sns.scatterplot( + data=df_v3, x='center_pca_normal_median_deg', y='v3_normal_median_deg', + hue='group', style='group', s=80, alpha=0.8, palette="deep" +) +# 画 y=x 的对角线 +lims = [10, 65] +plt.plot(lims, lims, '--', color='gray', zorder=0) +plt.fill_between(lims, lims, 70, color='red', alpha=0.05, zorder=0) # 上半区标红(V3更差) +plt.fill_between(lims, 0, lims, color='green', alpha=0.05, zorder=0) # 下半区标绿(V3更好) + +plt.title("Covariance Orientation vs. Center PCA Normal") +plt.xlabel("Center-PCA Normal Median Error (deg)") +plt.ylabel("V3 Covariance Normal Median Error (deg)") +plt.xlim(10, 65) +plt.ylim(25, 70) +plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', title="Method Group") +plt.text(12, 60, "V3 is Worse", color='red', fontsize=12, fontweight='bold', alpha=0.5) +plt.text(50, 30, "V3 is Better", color='green', fontsize=12, fontweight='bold', alpha=0.5) + +plt.tight_layout() +plt.savefig(OUT_DIR / "Fig1_V3_vs_PCA_Scatter.pdf", bbox_inches='tight') +print(f"✅ Saved Fig 1 to {OUT_DIR / 'Fig1_V3_vs_PCA_Scatter.pdf'}") + +# ========================================== +# Figure 2: Surface Near Ratio by Scene Barplot +# ========================================== +plt.figure(figsize=(7, 4)) +df_scene_sorted = df_scene.sort_values(by='surface_near_ratio', ascending=False) +sns.barplot(data=df_scene_sorted, x='scene', y='surface_near_ratio', palette="viridis") +plt.title("Proximity of Gaussian Centers to True Surface Across Scenes") +plt.xlabel("Tanks & Temples Scene") +plt.ylabel("Surface-Near Ratio ($< 2\\tau$)") +plt.axhline(y=df_scene['surface_near_ratio'].mean(), color='r', linestyle='--', label='Global Mean') +plt.legend() + +plt.tight_layout() +plt.savefig(OUT_DIR / "Fig2_Surface_Proximity_Bar.pdf") +print(f"✅ Saved Fig 2 to {OUT_DIR / 'Fig2_Surface_Proximity_Bar.pdf'}") diff --git a/scripts/plot_densification_timeseries.py b/scripts/plot_densification_timeseries.py new file mode 100644 index 0000000000000000000000000000000000000000..6af5a7da62fb0e551d71a6875e8507253fb80c72 --- /dev/null +++ b/scripts/plot_densification_timeseries.py @@ -0,0 +1,48 @@ +"""Plot edge/flat densification-event ratio over training steps (Exp 2 curve).""" +import os, json, argparse +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--methods", nargs="+", default=["vanilla_logged","vanillasgf_logged"]) + ap.add_argument("--scenes", nargs="+", default=["francis","m60","panther","bonsai","bicycle","garden"]) + ap.add_argument("--outputs_root", default="/root/autodl-tmp/SplatAtlas/outputs") + ap.add_argument("--out", default="/root/autodl-tmp/SplatAtlas/densification_timeseries.png") + ap.add_argument("--bin_size", type=int, default=1000) + args = ap.parse_args() + + fig, ax = plt.subplots(figsize=(8, 5)) + colors = {"vanilla_logged":"#888888", "vanillasgf_logged":"#1f77b4"} + for method in args.methods: + all_steps, all_ratios = [], [] + for scene in args.scenes: + p = os.path.join(args.outputs_root, f"{method}_{scene}", "densification_events.json") + if not os.path.exists(p): continue + for e in json.load(open(p)): + all_steps.append(e["step"]) + all_ratios.append(e["n_edge"] / max(e["n_flat"], 1)) + if not all_steps: + print(f"no data for {method}"); continue + steps, ratios = np.array(all_steps), np.array(all_ratios) + bins = np.arange(0, int(steps.max())+args.bin_size, args.bin_size) + bc, bm = [], [] + for i in range(len(bins)-1): + m = (steps >= bins[i]) & (steps < bins[i+1]) + if m.sum() > 0: + bc.append((bins[i]+bins[i+1])/2); bm.append(ratios[m].mean()) + ax.plot(bc, bm, label=method, marker='o', markersize=4, linewidth=1.5, + color=colors.get(method)) + ax.axhline(1.0, color="black", linewidth=0.5, linestyle="--", alpha=0.5) + ax.set_xlabel("Training step"); ax.set_ylabel("Edge / Flat densification-event ratio") + ax.legend(loc="best", frameon=False) + ax.set_title("Densification event redistribution over training") + plt.tight_layout(); plt.savefig(args.out, dpi=120, bbox_inches="tight") + print(f"saved {args.out}") + + +if __name__ == "__main__": + main() diff --git a/scripts/prepare_submit.py b/scripts/prepare_submit.py new file mode 100644 index 0000000000000000000000000000000000000000..b84b92a4fb96f016dcb950f34d8a9d5a14e9cb35 --- /dev/null +++ b/scripts/prepare_submit.py @@ -0,0 +1,40 @@ +import os +import zipfile +import argparse +import json + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--method", type=str, required=True) + parser.add_argument("--output_dir", type=str, default="outputs") + parser.add_argument("--zip_name", type=str, default=None) + return parser.parse_args() + +def main(): + args = parse_args() + zip_name = args.zip_name if args.zip_name else f"submission_{args.method}.zip" + + target_dirs = [d for d in os.listdir(args.output_dir) if d.startswith(f"{args.method}_")] + + if not target_dirs: + print(f"Error: No outputs found for method {args.method}") + return + + with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf: + for d in target_dirs: + base_path = os.path.join(args.output_dir, d) + metrics_path = os.path.join(base_path, "metrics.json") + if os.path.exists(metrics_path): + zipf.write(metrics_path, arcname=os.path.join(d, "metrics.json")) + + renders_path = os.path.join(base_path, "renders") + if os.path.exists(renders_path): + for img in os.listdir(renders_path): + if img.endswith(".png"): + img_path = os.path.join(renders_path, img) + zipf.write(img_path, arcname=os.path.join(d, "renders", img)) + + print(f"Submission package generated: {zip_name}") + +if __name__ == "__main__": + main() diff --git a/scripts/print_efficiency.py b/scripts/print_efficiency.py new file mode 100644 index 0000000000000000000000000000000000000000..3346aed33daef784f03cec0f12b3be3d7ce4aa77 --- /dev/null +++ b/scripts/print_efficiency.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +""" +Print Gaussian count, timing, and VRAM for all our completed experiments. + +Data sources (in priority order): + - Gaussian count : PLY header (element vertex N) + - iter_time_ms : TensorBoard train/iter_time_ms (last value) + - peak_vram_GB : TensorBoard train/peak_vram_GB (max across run) + - timing fallback : timing.json cumulative_gpu_time_sec / iters + +Usage: + python scripts/print_efficiency.py +""" +import os +import json +import glob +import numpy as np + +OUTPUTS_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +ITERS = 30000 + +KNOWN = ("vanillasgf", "vanilla_3dgs", "vanilla_logged", "vanilla_thresh", + "absgs", "pixelgs", "edgeloss") + +OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} +INDOOR_360 = {"bonsai", "counter", "kitchen", "room"} +TNT = { + "auditorium", "ballroom", "courtroom", "family", "francis", "horse", + "lighthouse", "m60", "museum", "palace", "panther", "playground", + "temple", "train" +} +ALL_SCENES = OUTDOOR_360 | INDOOR_360 | TNT + + +def is_ours(method): + return any(method == p or method.startswith(p) for p in KNOWN) + + +def parse_dir(d): + if d.startswith("vanilla_thresh_"): + rest = d[len("vanilla_thresh_"):] + for s in ALL_SCENES: + if rest == s or rest.startswith(s + "_"): + return "vanilla_thresh", s + return None, None + tok = d.split("_") + for i in range(len(tok) - 1, 0, -1): + sc = "_".join(tok[i:]) + me = "_".join(tok[:i]) + if sc in ALL_SCENES: + return me, sc + return None, None + + +def count_gaussians(ply_path): + try: + with open(ply_path, "rb") as f: + for line in f: + line = line.decode("ascii", errors="ignore").strip() + if line.startswith("element vertex"): + return int(line.split()[-1]) + if line == "end_header": + break + except Exception: + pass + return None + + +def read_tb(model_path): + """Read last iter_time_ms and max peak_vram_GB from TensorBoard events.""" + try: + from tensorboard.backend.event_processing.event_accumulator import EventAccumulator + ea = EventAccumulator(model_path, size_guidance={"scalars": 0}) + ea.Reload() + tags = ea.Tags().get("scalars", []) + + it_ms = None + for tag in ["train/iter_time_ms", "iter_time_ms"]: + if tag in tags: + vals = ea.Scalars(tag) + if vals: + # mean of last 10% steps (stable phase) + n = max(1, len(vals) // 10) + it_ms = float(np.mean([v.value for v in vals[-n:]])) + break + + vram = None + for tag in ["train/peak_vram_GB", "peak_vram_GB"]: + if tag in tags: + vals = ea.Scalars(tag) + if vals: + vram = float(max(v.value for v in vals)) + break + + return it_ms, vram + except Exception: + return None, None + + +def read_timing_json(model_path, iters): + """Fallback: cumulative_gpu_time_sec from timing.json.""" + try: + p = os.path.join(model_path, "timing.json") + d = json.load(open(p)) + cgt = d.get("cumulative_gpu_time_sec") + if cgt: + return float(cgt) / iters * 1000 # ms/iter + except Exception: + pass + return None + + +def main(): + plys = sorted(glob.glob( + f"{OUTPUTS_ROOT}/*/point_cloud/iteration_{ITERS}/point_cloud.ply")) + + rows = [] + for ply_path in plys: + parts = ply_path.split("/") + dirname = parts[-4] + method, scene = parse_dir(dirname) + if method is None or not is_ours(method): + continue + + model_path = os.path.join(OUTPUTS_ROOT, dirname) + n_gauss = count_gaussians(ply_path) + it_ms, vram = read_tb(model_path) + if it_ms is None: + it_ms = read_timing_json(model_path, ITERS) + + rows.append((dirname, method, scene, n_gauss, it_ms, vram)) + + # ── Print ────────────────────────────────────────────────────────────── + print(f"\n{'Dir':<42} {'#Gauss':>10} {'ms/iter':>9} {'PeakGB':>8}") + print("-" * 73) + + by_method = {} + for dirname, method, scene, n, t, v in rows: + ns = f"{n:,}" if n else "-" + ts = f"{t:.2f}" if t else "-" + vs = f"{v:.3f}" if v else "-" + print(f"{dirname:<42} {ns:>10} {ts:>9} {vs:>8}") + by_method.setdefault(method, []).append((n, t, v)) + + # ── Per-method averages ──────────────────────────────────────────────── + print(f"\n{'Method':<35} {'Avg #Gauss':>11} {'Avg ms/iter':>12} {'Avg PeakGB':>11} {'N':>4}") + print("-" * 67) + for m, vals in sorted(by_method.items()): + ns = [x[0] for x in vals if x[0] is not None] + ts = [x[1] for x in vals if x[1] is not None] + vs = [x[2] for x in vals if x[2] is not None] + nav = f"{int(np.mean(ns)):,}" if ns else "-" + tav = f"{np.mean(ts):.2f}" if ts else "-" + vav = f"{np.mean(vs):.3f}" if vs else "-" + print(f"{m:<35} {nav:>11} {tav:>12} {vav:>11} {len(vals):>4}") + + +if __name__ == "__main__": + main() diff --git a/scripts/region_psnr_analysis.py b/scripts/region_psnr_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..26b06a9db3637474d31f1cbe0f1a2214d936919e --- /dev/null +++ b/scripts/region_psnr_analysis.py @@ -0,0 +1,67 @@ +"""Region-level PSNR (edge/flat by Sobel quartile of GT) → Table 7 expansion.""" +import os, argparse, glob +import numpy as np +from PIL import Image +import torch +import torch.nn.functional as F + + +def sobel_mag(img_chw): + luma = (0.299*img_chw[0] + 0.587*img_chw[1] + 0.114*img_chw[2]).unsqueeze(0).unsqueeze(0) + wx = torch.tensor([[1,0,-1],[2,0,-2],[1,0,-1]], dtype=torch.float32).view(1,1,3,3) + wy = torch.tensor([[1,2,1],[0,0,0],[-1,-2,-1]], dtype=torch.float32).view(1,1,3,3) + gx = F.conv2d(luma, wx, padding=1) + gy = F.conv2d(luma, wy, padding=1) + return torch.sqrt(gx**2 + gy**2 + 1e-8).squeeze() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--methods", nargs="+", default=["vanilla_logged", "vanillasgf_logged", + "vanillasgf_nosg", "vanillasgf_hook"]) + ap.add_argument("--scenes", nargs="+", default=["francis","m60","panther","bonsai","bicycle","garden"]) + ap.add_argument("--outputs_root", default="/root/autodl-tmp/SplatAtlas/outputs") + ap.add_argument("--iter", type=int, default=30000) + args = ap.parse_args() + + print(f"{'Method':<22} {'Scene':<10} {'Edge PSNR':>10} {'Flat PSNR':>10} {'Full PSNR':>10}") + print("-"*68) + for method in args.methods: + all_e, all_f, all_t = [], [], [] + for scene in args.scenes: + rd = os.path.join(args.outputs_root, f"{method}_{scene}", f"renders_test_{args.iter}") + gd = os.path.join(args.outputs_root, f"{method}_{scene}", f"gt_test_{args.iter}") + if not (os.path.exists(rd) and os.path.exists(gd)): + print(f"{method:<22} {scene:<10} (renders/gt missing)"); continue + files = sorted(glob.glob(os.path.join(rd, "*.png")) + glob.glob(os.path.join(rd, "*.jpg"))) + se, sf, st = [], [], [] + for pf in files: + gf = os.path.join(gd, os.path.basename(pf)) + if not os.path.exists(gf): continue + pred = torch.from_numpy(np.array(Image.open(pf).convert("RGB")).astype(np.float32)/255.).permute(2,0,1) + gt = torch.from_numpy(np.array(Image.open(gf).convert("RGB")).astype(np.float32)/255.).permute(2,0,1) + if pred.shape != gt.shape: continue + sob = sobel_mag(gt) + q75, q25 = float(sob.quantile(0.75)), float(sob.quantile(0.25)) + em, fm = (sob >= q75), (sob <= q25) + d = ((pred-gt)**2).mean(0) + se.append(float(d[em].mean()) if em.any() else 0) + sf.append(float(d[fm].mean()) if fm.any() else 0) + st.append(float(d.mean())) + if not st: + print(f"{method:<22} {scene:<10} (no images)"); continue + ep = -10*np.log10(np.mean(se)+1e-12) + fp = -10*np.log10(np.mean(sf)+1e-12) + tp = -10*np.log10(np.mean(st)+1e-12) + print(f"{method:<22} {scene:<10} {ep:>10.2f} {fp:>10.2f} {tp:>10.2f}") + all_e.extend(se); all_f.extend(sf); all_t.extend(st) + if all_t: + print(f"{method:<22} {'AVERAGE':<10} " + f"{-10*np.log10(np.mean(all_e)+1e-12):>10.2f} " + f"{-10*np.log10(np.mean(all_f)+1e-12):>10.2f} " + f"{-10*np.log10(np.mean(all_t)+1e-12):>10.2f}") + print("-"*68) + + +if __name__ == "__main__": + main() diff --git a/scripts/repair_seed_variance.py b/scripts/repair_seed_variance.py new file mode 100644 index 0000000000000000000000000000000000000000..3490113d78b40ff3e3bb9853a9aa1562cb34d876 --- /dev/null +++ b/scripts/repair_seed_variance.py @@ -0,0 +1,273 @@ +#!/usr/bin/env python3 +""" +repair_seed_variance.py +针对 5 个 CRITICAL 组的定点修复脚本 +运行方式: python /root/autodl-tmp/SplatAtlas/scripts/repair_seed_variance.py +""" + +import os, subprocess, shutil, json, time +from pathlib import Path + +SPLATATLAS = "/root/autodl-tmp/SplatAtlas" +OUTPUTS = f"{SPLATATLAS}/outputs" +LOGS = f"{SPLATATLAS}/logs/repair" +os.makedirs(LOGS, exist_ok=True) + +ITERS = 30000 +LOG_FILE = open(f"{LOGS}/repair_{time.strftime('%Y%m%d_%H%M%S')}.log", "a") + +def log(msg): + ts = time.strftime("%H:%M:%S") + line = f"[{ts}] {msg}" + print(line); LOG_FILE.write(line + "\n"); LOG_FILE.flush() + +def run(cmd, env=None, cwd=SPLATATLAS): + log(f" CMD: {cmd[:120]}...") + r = subprocess.run(cmd, shell=True, env=env, cwd=cwd) + ok = r.returncode == 0 + log(f" -> {'OK' if ok else 'FAILED'} (rc={r.returncode})") + return ok + +def make_env(env_path, extra_py, oom_safe=False): + """构造运行环境：PYTHONPATH + 可选 OOM 防护""" + e = os.environ.copy() + e["PYTHONPATH"] = f"{extra_py}:{SPLATATLAS}:{e.get('PYTHONPATH', '')}" + if oom_safe: + e["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128" + return e + +def python_bin(env_path): + return "python" if env_path == "base" else f"{env_path}/bin/python" + +def check_ply(model_dir): + p = f"{model_dir}/point_cloud/iteration_{ITERS}/point_cloud.ply" + return os.path.exists(p), p + +def check_flag(model_dir): + p = f"{model_dir}/render_complete_{ITERS}.flag" + return os.path.exists(p) + +def check_metrics(model_dir): + p = f"{model_dir}/metrics_test_iter{ITERS}.json" + return os.path.exists(p), p + +def do_train(python, env, method, source, model_dir, res, seed=None): + seed_str = f" --seed {seed}" if seed else "" + cmd = ( + f"{python} {SPLATATLAS}/scripts/main_train.py" + f" --method {method}" + f" --source_path {source}" + f" --model_path {model_dir}" + f" --iterations {ITERS}" + f" --resolution {res}" + f" --save_iterations 30000" + f"{seed_str}" + ) + return run(cmd, env) + +def do_render(python, env, method, source, model_dir, res): + # 清理残留渲染目录 + for d in [f"renders_test_{ITERS}", f"gt_test_{ITERS}", f"depths_test_{ITERS}"]: + p = f"{model_dir}/{d}" + if os.path.exists(p): shutil.rmtree(p) + cmd = ( + f"{python} {SPLATATLAS}/scripts/main_render.py" + f" --method {method}" + f" --source_path {source}" + f" --model_path {model_dir}" + f" --iteration {ITERS}" + f" --resolution {res}" + ) + return run(cmd, env) + +def do_eval(python, env, method, scene, model_dir, source): + ply_path = f"{model_dir}/point_cloud/iteration_{ITERS}/point_cloud.ply" + render_dir = f"{model_dir}/renders_test_{ITERS}" + gt_dir = f"{model_dir}/gt_test_{ITERS}" + depth_dir = f"{model_dir}/depths_test_{ITERS}" + out_json = f"{model_dir}/metrics_test_iter{ITERS}.json" + cmd = ( + f"{python} {SPLATATLAS}/ufd_evalkit/run_eval.py" + f" --method {method}" + f" --scene {scene}" + f" --render_dir {render_dir}" + f" --gt_dir {gt_dir}" + f" --ply_path {ply_path}" + f" --output_json {out_json}" + f" --colmap_dir {source}" + f" --depth_dir {depth_dir}" + ) + return run(cmd, env) + +# ───────────────────────────────────────────────────────────────────────────── +# 修复目标定义 +# ───────────────────────────────────────────────────────────────────────────── + +TARGETS = [ + + # ── Case 1: vanilla_3dgs bonsai ────────────────────────────────────────── + # PLY 存在，渲染存在，只是 eval OOM。加 OOM 防护重跑 eval。 + { + "label": "vanilla_3dgs bonsai (all seeds, eval OOM fix)", + "method": "vanilla_3dgs", + "scene": "bonsai", + "source": "/root/autodl-tmp/dataset/360/bonsai", + "res": 2, + "env_path": "base", + "extra_py": "/root/autodl-tmp/3dgs_official", + "oom_safe": True, # ← 关键：防 LPIPS OOM + "seeds": [None, 1, 2], # None = default (无 seed suffix) + }, + + # ── Case 2: analyticsplatting bonsai seed1/seed2 ───────────────────────── + # default 已有 metrics，seed1/seed2 缺 PLY → 需重训 + { + "label": "analyticsplatting bonsai (seed1, seed2 retrain)", + "method": "analyticsplatting", + "scene": "bonsai", + "source": "/root/autodl-tmp/dataset/360/bonsai", + "res": 2, + "env_path": "/root/autodl-tmp/envs/ana_splatting", + "extra_py": "/root/autodl-tmp/Analytic-Splatting_offy", + "oom_safe": False, + "seeds": [1, 2], # default 已完成，只补 seed1/seed2 + }, + + # ── Case 3: analyticsplatting Lego seed1/seed2 ─────────────────────────── + { + "label": "analyticsplatting Lego (seed1, seed2 retrain)", + "method": "analyticsplatting", + "scene": "Lego", + "source": "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/Lego", + "res": 1, + "env_path": "/root/autodl-tmp/envs/ana_splatting", + "extra_py": "/root/autodl-tmp/Analytic-Splatting_offy", + "oom_safe": False, + "seeds": [1, 2], + }, + + # ── Case 4: erankgs bonsai ─────────────────────────────────────────────── + # 所有 seed 都缺 metrics；render 崩溃是因为环境混用。 + # 用 erank_gs 环境严格跑 train→render→eval 全流程 + { + "label": "erankgs bonsai (all seeds, env fix)", + "method": "erankgs", + "scene": "bonsai", + "source": "/root/autodl-tmp/dataset/360/bonsai", + "res": 2, + "env_path": "/root/autodl-tmp/envs/erank_gs", + "extra_py": "/root/autodl-tmp/erank_gs", + "oom_safe": False, + "seeds": [None, 1, 2], + }, + + # ── Case 5: pgsr bonsai ────────────────────────────────────────────────── + # 与 erankgs 症状相同，严格用 pgsr_v2 环境 + { + "label": "pgsr bonsai (all seeds, env fix)", + "method": "pgsr", + "scene": "bonsai", + "source": "/root/autodl-tmp/dataset/360/bonsai", + "res": 2, + "env_path": "/root/autodl-tmp/envs/pgsr_v2", + "extra_py": "/root/autodl-tmp/pgsr_official", + "oom_safe": False, + "seeds": [None, 1, 2], + }, +] + +# ───────────────────────────────────────────────────────────────────────────── +# 执行逻辑 +# ───────────────────────────────────────────────────────────────────────────── + +def model_dir_for(method, scene, seed): + """None seed → vanilla 目录名（无 suffix）""" + suffix = f"_seed{seed}" if seed is not None else "" + return f"{OUTPUTS}/{method}{suffix}_{scene}" + +def repair(target): + method = target["method"] + scene = target["scene"] + source = target["source"] + res = target["res"] + env_path = target["env_path"] + extra_py = target["extra_py"] + oom_safe = target["oom_safe"] + python = python_bin(env_path) + env = make_env(env_path, extra_py, oom_safe) + + log(f"\n{'='*60}") + log(f"TARGET: {target['label']}") + + for seed in target["seeds"]: + mdir = model_dir_for(method, scene, seed) + tag = f"{method}_seed{seed}_{scene}" if seed else f"{method}_{scene}" + log(f"\n -- {tag} --") + log(f" dir: {mdir}") + + # ── 1. 诊断当前状态 ────────────────────────────────────────────────── + ply_ok, ply_path = check_ply(mdir) + flag_ok = check_flag(mdir) + metrics_ok, _ = check_metrics(mdir) + + if metrics_ok: + log(" metrics 已存在，跳过") + continue + + log(f" 状态: PLY={'OK' if ply_ok else 'MISSING'} " + f"render={'OK' if flag_ok else 'MISSING'} " + f"metrics=MISSING") + + # ── 2. 训练（如需）────────────────────────────────────────────────── + if not ply_ok: + log(" → 执行训练") + if not do_train(python, env, method, source, mdir, res, seed): + log(f" !! 训练失败，跳过后续步骤: {tag}") + continue + ply_ok = True + flag_ok = False # 训练完需重新渲染 + + # ── 3. 渲染（如需）────────────────────────────────────────────────── + if not flag_ok: + log(" → 执行渲染") + if not do_render(python, env, method, source, mdir, res): + log(f" !! 渲染失败，跳过 eval: {tag}") + continue + + # ── 4. 评估 ───────────────────────────────────────────────────────── + log(" → 执行评估") + if do_eval(python, env, method, scene, mdir, source): + log(f" ✓ DONE: {tag}") + else: + log(f" !! eval 失败: {tag}") + +# ───────────────────────────────────────────────────────────────────────────── +# 收尾：汇总状态 +# ───────────────────────────────────────────────────────────────────────────── + +def final_report(): + log("\n" + "="*60) + log("修复后状态汇总") + log("="*60) + header = f"{'Group':<40} {'PLY':>5} {'Metrics':>8}" + log(header) + for t in TARGETS: + for seed in t["seeds"]: + mdir = model_dir_for(t["method"], t["scene"], seed) + tag = f"{t['method']}_seed{seed}_{t['scene']}" if seed else f"{t['method']}_{t['scene']}" + ply_ok,_ = check_ply(mdir) + metrics_ok,_ = check_metrics(mdir) + log(f" {tag:<38} {'✓' if ply_ok else '✗':>5} {'✓' if metrics_ok else '✗':>8}") + +if __name__ == "__main__": + log("====== repair_seed_variance 开始 ======") + t0 = time.time() + + for target in TARGETS: + repair(target) + + final_report() + + elapsed = (time.time() - t0) / 3600 + log(f"\n====== 完成，总耗时 {elapsed:.1f}h ======") + LOG_FILE.close() \ No newline at end of file diff --git a/scripts/rescue_eval.py b/scripts/rescue_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..e87db8b28383e3c46f349974a8a3d19b4d06dd2f --- /dev/null +++ b/scripts/rescue_eval.py @@ -0,0 +1,80 @@ +import os +import subprocess +import yaml +import glob + +def build_scene_map(): + """反向解析 configs 目录下的所有 yaml，提取真实的源路径映射""" + scene_map = {} + config_files = glob.glob("/root/autodl-tmp/SplatAtlas/configs/*.yaml") + for cfg_file in config_files: + try: + with open(cfg_file, 'r') as f: + cfg = yaml.safe_load(f) + if "datasets" in cfg: + for ds in cfg["datasets"]: + base = ds.get("base_path", "") + for sc in ds.get("scenes", []): + scene_map[sc] = os.path.join(base, sc) + except Exception: + pass + return scene_map + +def main(): + base_dir = "/root/autodl-tmp/SplatAtlas/outputs" + iteration = 30000 + + # 建立 Scene -> Source Path 的绝对真理映射 + scene_map = build_scene_map() + + for folder in os.listdir(base_dir): + if not folder.startswith("2dgs_"): continue + + m_path = os.path.join(base_dir, folder) + method, scene = folder.split('_', 1) + + # 精准获取源路径，增加后备启发式搜索以防万一 + src = scene_map.get(scene, "") + if not src or not os.path.exists(src): + roots = [ + "/root/autodl-tmp/dataset/synthetic_nerf/Synthetic_NeRF", + "/root/autodl-tmp/dataset/mipnerf360", + "/root/autodl-tmp/dataset/mipnerf360/360_v2", # Mip-NeRF 360 的常见嵌套 + "/root/autodl-tmp/dataset/deepblending_clean" + ] + for r in roots: + p = os.path.join(r, scene) + if os.path.exists(p): + src = p + break + + if not src or not os.path.exists(src): + print(f"❌ [Fatal] 彻底找不到场景 {scene} 的数据集目录，跳过。") + continue + + # 补测基础评估 + metrics_json = os.path.join(m_path, f"metrics_test_iter{iteration}.json") + if not os.path.exists(metrics_json): + print(f"🛠️ [Eval] 正在抢救基础评测: {folder}") + cmd = (f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py " + f"--method {method} --scene {scene} " + f"--render_dir {m_path}/renders_test_{iteration} " + f"--gt_dir {m_path}/gt_test_{iteration} " + f"--ply_path {m_path}/point_cloud/iteration_{iteration}/point_cloud.ply " + f"--output_json {metrics_json} --colmap_dir {src} --skip_quarantine") + subprocess.run(cmd, shell=True) + + # 补测物理诊断 + physics_json = os.path.join(m_path, f"offline_physics_{iteration}.json") + if not os.path.exists(physics_json): + print(f"🛠️ [Physics] 正在抢救物理诊断: {folder}") + cmd = (f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py " + f"--method {method} --source_path {src} --model_path {m_path} --iteration {iteration}") + subprocess.run(cmd, shell=True) + + # 聚合大榜单 + print("📈 [GrandTable] 正在重构汇总榜单...") + subprocess.run("python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", shell=True) + +if __name__ == "__main__": + main() diff --git a/scripts/run_2dgs_benchmark.py b/scripts/run_2dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..7a07386199ae08a069226abd823e7a7b3ba89570 --- /dev/null +++ b/scripts/run_2dgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/2dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/2dgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"Done: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_2dgssgf_tnt.py b/scripts/run_2dgssgf_tnt.py new file mode 100644 index 0000000000000000000000000000000000000000..9a98d0e67b742230eab310429850043c71ecd612 --- /dev/null +++ b/scripts/run_2dgssgf_tnt.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/2dgssgf_tnt_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/2dgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"Done: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_3dgs_dr_benchmark.py b/scripts/run_3dgs_dr_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..ee825b2b82bdcdc40b4c17f4b3355d9d716399f4 --- /dev/null +++ b/scripts/run_3dgs_dr_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/3dgs_dr_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/3DGS-DR:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_3dgsmcmc_benchmark.py b/scripts/run_3dgsmcmc_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..19b8eb9c92bfb710be569c5ba5787ff8d9bd35c0 --- /dev/null +++ b/scripts/run_3dgsmcmc_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f"[{name}] Starting...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f"[{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f"[{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f"[{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/3dgsmcmc_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/3dgsmcmc_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print("[Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print("[Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print("[Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training and In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print("[Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_3dgut_benchmark.py b/scripts/run_3dgut_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..77910c5107f14e21c754bed74b5283ae8a413da8 --- /dev/null +++ b/scripts/run_3dgut_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [x] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [*] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [x] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/3dgut_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/3dgrut:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[*] [{method.upper()}] Processing Scene: {scene} (r={res})") + + ckpt_path = os.path.join(model_out, f"ours_{iters}/ckpt_{iters}.pt") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ckpt_path) or os.path.exists(os.path.join(model_out, "ckpt_last.pt")): + print(f" [*] [Checkpoint] Weights found.") + if not os.path.exists(render_flag): + print(f" [!] [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [*] [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [*] [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[*] {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_IJCV.sh b/scripts/run_IJCV.sh new file mode 100644 index 0000000000000000000000000000000000000000..b6faedeaabdca3b6258e8977093f52c88422a3dd --- /dev/null +++ b/scripts/run_IJCV.sh @@ -0,0 +1,126 @@ +#!/bin/bash +# Run 17 verified-by-smoke methods. +# Confidence levels noted inline. Failures of one method do NOT abort the rest. + +source /root/miniconda3/etc/profile.d/conda.sh +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +ERROR_LOG="/root/autodl-tmp/SplatAtlas/run_errors.log" +STATUS_LOG="/root/autodl-tmp/SplatAtlas/run_status.log" + +echo "================================================================" | tee -a "$ERROR_LOG" "$STATUS_LOG" +echo "17-method run started at $(date)" | tee -a "$ERROR_LOG" "$STATUS_LOG" +echo "================================================================" | tee -a "$ERROR_LOG" "$STATUS_LOG" + +# ---------------------------------------------------------------- +# Pre-flight: lp_3dgs env is missing torch (audit confirmed). +# Install before queue starts so it doesn't waste a slot failing on import. +# ---------------------------------------------------------------- +echo "[pre-flight] checking lp_3dgs torch..." +if ! /root/autodl-tmp/envs/lp_3dgs/bin/python -c "import torch" 2>/dev/null; then + echo "[pre-flight] lp_3dgs missing torch, installing torch==2.1.2+cu118..." + conda activate /root/autodl-tmp/envs/lp_3dgs + pip install torch==2.1.2 torchvision torchaudio \ + --index-url https://download.pytorch.org/whl/cu118 -q + if /root/autodl-tmp/envs/lp_3dgs/bin/python -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then + echo "[pre-flight] lp_3dgs torch OK" + else + echo "[pre-flight] lp_3dgs torch install FAILED — will skip lp_3dgs in queue" \ + | tee -a "$ERROR_LOG" + SKIP_LP3DGS=1 + fi + conda deactivate +fi + +# ---------------------------------------------------------------- +# Task list. Format: " # " +# Order: low risk → high risk +# ---------------------------------------------------------------- +declare -a tasks=( + # --- Tier 1: smoke 3000-iter user-confirmed exit 0 --- + "/root/autodl-tmp/envs/gigs scripts/run_gigs_benchmark.py" + "/root/autodl-tmp/envs/sort_free_gs scripts/run_sortfreegs_benchmark.py" + "/root/autodl-tmp/envs/analytic_splatting scripts/run_analyticsplatting_benchmark.py" + "/root/autodl-tmp/envs/ngs scripts/run_ngs_benchmark.py" + "/root/autodl-tmp/envs/aaa_gaussians scripts/run_aaags_benchmark.py" + "/root/autodl-tmp/envs/improving_adc scripts/run_improvingadc_benchmark.py" + "/root/autodl-tmp/envs/compact3dgs scripts/run_compact3dgs_benchmark.py" + "/root/autodl-tmp/envs/sogs scripts/run_sogs_benchmark.py" + + # --- Tier 2: 100-iter audit OK (worker B 报告过 + audit 复核) --- + "/root/autodl-tmp/envs/cdc_gs scripts/run_cdcgs_benchmark.py" + "/root/autodl-tmp/envs/op43dgs scripts/run_op43dgs_benchmark.py" + "/root/autodl-tmp/envs/bags scripts/run_bags_benchmark.py" + "/root/autodl-tmp/envs/gspull scripts/run_gspull_benchmark.py" + "/root/autodl-tmp/envs/FLoD scripts/run_flod_benchmark.py" + + # --- Tier 3: 慢 / 弱信号(silent failure 风险),放后面 --- + "/root/autodl-tmp/envs/spec_gaussian scripts/run_specgaussian_benchmark.py" + "/root/autodl-tmp/envs/pup_3dgs scripts/run_pup3dgs_benchmark.py" + "/root/autodl-tmp/envs/gaussianfocus scripts/run_gaussianfocus_benchmark.py" + "/root/autodl-tmp/envs/lp_3dgs scripts/run_lp3dgs_benchmark.py" +) + +# ---------------------------------------------------------------- +# Run loop +# ---------------------------------------------------------------- +TOTAL=${#tasks[@]} +IDX=0 +for task in "${tasks[@]}"; do + IDX=$((IDX + 1)) + read -r env_path script_path <<< "$task" + method_name=$(basename "$script_path" .py) + + echo "----------------------------------------------------------------" \ + | tee -a "$STATUS_LOG" + echo "[$IDX/$TOTAL] $(date '+%Y-%m-%d %H:%M:%S') $method_name" \ + | tee -a "$STATUS_LOG" + + # Skip lp_3dgs if pre-flight torch install failed + if [ "$method_name" = "run_lp3dgs_benchmark" ] && [ "$SKIP_LP3DGS" = "1" ]; then + echo "[skip] lp_3dgs torch not available" | tee -a "$STATUS_LOG" "$ERROR_LOG" + continue + fi + + # Sanity: env exists? + if [ ! -d "$env_path" ]; then + echo "[skip] env missing: $env_path" | tee -a "$STATUS_LOG" "$ERROR_LOG" + continue + fi + + conda activate "$env_path" + if [ $? -ne 0 ]; then + echo "[fail] could not activate $env_path" | tee -a "$STATUS_LOG" "$ERROR_LOG" + continue + fi + + # Common deps (idempotent, fast if already there) + pip install lpips imageio scikit-image -q 2>/dev/null + + START=$(date +%s) + PYTHONPATH=/root/autodl-tmp/SplatAtlas \ + "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" \ + 2> >(tee -a "$ERROR_LOG" >&2) + RC=$? + END=$(date +%s) + DURATION=$((END - START)) + + if [ $RC -eq 0 ]; then + echo "[done] $method_name (rc=0, ${DURATION}s)" | tee -a "$STATUS_LOG" + else + echo "[FAIL] $method_name (rc=$RC, ${DURATION}s) — see $ERROR_LOG" \ + | tee -a "$STATUS_LOG" + echo "[$(date)] FAILED: $method_name (rc=$RC)" >> "$ERROR_LOG" + fi + + conda deactivate +done + +echo "================================================================" | tee -a "$STATUS_LOG" +echo "17-method run finished at $(date)" | tee -a "$STATUS_LOG" +echo "Status log: $STATUS_LOG" +echo "Error log: $ERROR_LOG" +echo "" +echo "To check render flags after run:" +echo " ls outputs/*/render_complete_30000.flag 2>/dev/null | wc -l" +echo "================================================================" \ No newline at end of file diff --git a/scripts/run_aaags_benchmark.py b/scripts/run_aaags_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..7a7fb0ae1aefc4bde89f509c28ca8b2a13cde112 --- /dev/null +++ b/scripts/run_aaags_benchmark.py @@ -0,0 +1,66 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + return False + return True + except Exception as e: + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/aaags_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/AAAGaussians:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_absgs_benchmark.py b/scripts/run_absgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..3f4e2a27e56e396faadec5f0d12ba4cf6dde1977 --- /dev/null +++ b/scripts/run_absgs_benchmark.py @@ -0,0 +1,74 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/absgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/absgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_absgssgf.py b/scripts/run_absgssgf.py new file mode 100644 index 0000000000000000000000000000000000000000..4d28861281c0e7e3281fe8e309af5c3065a9955e --- /dev/null +++ b/scripts/run_absgssgf.py @@ -0,0 +1,83 @@ +import os +import yaml +import subprocess +import time +import shutil +import glob +import re + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/absgssgf_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/absgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + ckpt_list = glob.glob(os.path.join(model_out, "chkpnt*.pth")) + if ckpt_list: + latest_ckpt = max(ckpt_list, key=lambda x: int(re.findall(r'chkpnt(\d+)\.pth', os.path.basename(x))[0] if re.findall(r'chkpnt(\d+)\.pth', os.path.basename(x)) else 0)) + print(f" [Resume] 发现断点，从 {latest_ckpt} 续训...") + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling --start_checkpoint {latest_ckpt}" + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_all_42_methods.sh b/scripts/run_all_42_methods.sh new file mode 100644 index 0000000000000000000000000000000000000000..6465a15cc8d9057a3cb8c283d920ab4bf327dcc5 --- /dev/null +++ b/scripts/run_all_42_methods.sh @@ -0,0 +1,116 @@ +#!/bin/bash + +# 初始化 Conda 环境路径 +source /root/miniconda3/etc/profile.d/conda.sh + +# 确保主路径在 PYTHONPATH 中 +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +# 定义错误日志路径 +ERROR_LOG="/root/autodl-tmp/SplatAtlas/run_errors.log" + +# 每次运行前，在错误日志里画个分割线 +echo "================================================================" >> "$ERROR_LOG" +echo "🚀 New Run Started at $(date)" >> "$ERROR_LOG" +echo "================================================================" >> "$ERROR_LOG" + +# 定义任务队列：(环境绝对路径 脚本相对路径) +declare -a tasks=( + # --- 一、 Initialization (2) --- + "base scripts/run_vanilla_benchmark.py" + "/root/autodl-tmp/envs/3dgs_mcmc scripts/run_3dgsmcmc_benchmark.py" + +# --- 二、 Attribute Expansion (10) --- + "/root/autodl-tmp/envs/2dgs_v2 scripts/run_2dgs_benchmark.py" + "/root/autodl-tmp/envs/pgsr scripts/run_pgsr_benchmark.py" + "/root/autodl-tmp/envs/ges scripts/run_ges_benchmark.py" + "/root/autodl-tmp/envs/scaffold_gs scripts/run_scaffoldgs_benchmark.py" + "/root/autodl-tmp/envs/r3dg scripts/run_r3dg_benchmark.py" + "/root/autodl-tmp/envs/3dgrut scripts/run_3dgut_benchmark.py" + "/root/autodl-tmp/envs/hogs scripts/run_hogs_benchmark.py" + "/root/autodl-tmp/envs/texture_gs scripts/run_texturegs_benchmark.py" + "/root/autodl-tmp/envs/textured_gs scripts/run_texturedgs_benchmark.py" + "/root/autodl-tmp/envs/gaussian_surfels scripts/run_surfel_benchmark.py" + + # --- 三、 Rendering & Rasterization (5) --- + "/root/autodl-tmp/envs/mip-splatting scripts/run_mipsplatting_benchmark.py" + "/root/autodl-tmp/envs/ana_splatting scripts/run_analyticsplatting_benchmark.py" + "/root/autodl-tmp/envs/3dgs_dr scripts/run_3dgs_dr_benchmark.py" + "/root/autodl-tmp/envs/lod_gs scripts/run_lod_gs_benchmark.py" + "/root/autodl-tmp/envs/octree_gs scripts/run_octree_gs_benchmark.py" + + # --- 四、 Regularization (9) --- + "/root/autodl-tmp/env scripts/run_sugar_benchmark.py" + "/root/autodl-tmp/envs/gof scripts/run_gof_benchmark.py" + "/root/autodl-tmp/envs/gsdf scripts/run_gsdf_benchmark.py" + "/root/autodl-tmp/envs/erank_gs scripts/run_erankgs_benchmark.py" + "/root/autodl-tmp/envs/trimgs scripts/run_trimgs_benchmark.py" + "/root/autodl-tmp/envs/gaussianpro scripts/run_gaussianpro_benchmark.py" + "/root/autodl-tmp/envs/gspull scripts/run_gspull_benchmark.py" + "/root/autodl-tmp/envs/lapis_gs scripts/run_lapisgs_benchmark.py" + + + # --- 五、 Training Strategy (2) --- + "/root/autodl-tmp/envs/opti3dgs scripts/run_opti3dgs_benchmark.py" + "/root/autodl-tmp/envs/co_adaptation_3dgs scripts/run_coadaptgs_benchmark.py" + + # --- 六、 Adaptive Control (10) --- + "/root/autodl-tmp/envs/absgs_env scripts/run_absgs_benchmark.py" + "/root/autodl-tmp/envs/pixelgs scripts/run_pixelgs_benchmark.py" + "/root/autodl-tmp/envs/mini_splatting scripts/run_minisplatting_benchmark.py" + "/root/autodl-tmp/envs/taming_3dgs scripts/run_tamgs_benchmark.py" + "/root/autodl-tmp/envs/atomgs scripts/run_atomgs_benchmark.py" + "/root/autodl-tmp/envs/reactgs_env scripts/run_reactgs_benchmark.py" + "/root/autodl-tmp/envs/gs_lpm scripts/run_gslpm.py" + "/root/autodl-tmp/envs/conegs_env scripts/run_conegs_benchmark.py" + "/root/autodl-tmp/envs/steepgs_env scripts/run_steepgs.py" + + # --- 七、 Post-Processing (4) --- + "/root/autodl-tmp/envs/eagles scripts/run_eagles_benchmark.py" + "/root/autodl-tmp/envs/rdo-gaussian scripts/run_rdogaussian_benchmark.py" + "/root/autodl-tmp/envs/ghap_env scripts/run_ghap_benchmark.py" + "/root/autodl-tmp/envs/hac_plus scripts/run_hac_plus_benchmark.py" + "/root/autodl-tmp/envs/lightgaussian_official scripts/run_lightgaussian_benchmark.py" + ) + +echo "🚀 SplatAtlas: Starting Robust Optimization Pipeline (42 Methods)..." +echo "📝 Errors will be logged to: $ERROR_LOG" + +for task in "${tasks[@]}"; do + # 解析环境路径和脚本 + read -r env_path script_path <<< "$task" + method_name=$(basename "$script_path") + + echo "----------------------------------------------------------------" + echo "🔄 Method: $method_name" + + # 1. 激活环境 + conda activate "$env_path" + + if [ $? -eq 0 ]; then + # 2. 环境自动体检 + echo "🛠️ Checking requirements in $env_path..." + pip install lpips imageio scikit-image -q + + # 3. 执行 Benchmark，stderr 同时打印到终端并写入日志 + echo "✅ Environment ready. Executing..." + + "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" 2> >(tee -a "$ERROR_LOG" >&2) + + if [ $? -eq 0 ]; then + echo "✨ Finished: $method_name" + else + echo "❌ Error: $method_name failed. Saved to $ERROR_LOG" + echo "[$(date)] FAILED: $method_name" >> "$ERROR_LOG" + fi + else + echo "❌ Failed to activate: $env_path" + echo "[$(date)] FAILED (Env Activation): $env_path for $method_name" >> "$ERROR_LOG" + fi + + # 释放当前环境 + conda deactivate +done + +echo "================================================================" +echo "🏁 Optimization Pipeline completed. Check $ERROR_LOG for crash reports." \ No newline at end of file diff --git a/scripts/run_all_benchmarks_1.sh b/scripts/run_all_benchmarks_1.sh new file mode 100644 index 0000000000000000000000000000000000000000..661942a38738661e3f59ccc7ad43ad500bf07b2d --- /dev/null +++ b/scripts/run_all_benchmarks_1.sh @@ -0,0 +1,90 @@ +#!/bin/bash + +# 初始化 Conda 环境路径 +CONDA_BASE=$(conda info --base) +source "$CONDA_BASE/etc/profile.d/conda.sh" + +# 确保主路径在 PYTHONPATH 中，防止 eval 脚本找不到模块 +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +# 定义任务队列：(环境绝对路径 脚本相对路径) +declare -a tasks=( + # --- 分支 1.1：降维表征与面元化 --- + + + + + + + + # --- 分支 1.4：显式网格提取与拓扑绑定 --- + # "/root/autodl-tmp/envs/cadgs scripts/run_bgtriangle_benchmark.py" + + # --- 分支 1.5：复杂光照与材质的几何解耦 --- + # "/root/autodl-tmp/envs/3dgrut scripts/run_3dgut_benchmark.py" + # "/root/autodl-tmp/envs/gaussian_shader scripts/run_gaussianshader.py" + + # --- 分支 1.6：无界场景的空间映射与拓扑收缩 --- + # "/root/autodl-tmp/envs/hogs scripts/run_hogs_benchmark.py" + # "/root/autodl-tmp/envs/r3dg scripts/run_r3dg_benchmark.py" + + # --- 分支 1.7：结构化骨架与混合几何表征 --- + # "/root/autodl-tmp/envs/scaffold_gs scripts/run_scaffoldgs_benchmark.py" + + # --- 分支 1.8：高频特征与微观拓扑逼近 --- + "/root/autodl-tmp/envs/ges scripts/run_ges_benchmark.py" + # "/root/autodl-tmp/envs/texture_gs scripts/run_texturegs_benchmark.py" + "/root/autodl-tmp/envs/textured_gs scripts/run_texturedgs_benchmark.py" + + # "/root/autodl-tmp/envs/2dgs scripts/run_2dgs_benchmark.py" + "/root/autodl-tmp/envs/gaussian_surfels scripts/run_surfel_benchmark.py" + + + "/root/autodl-tmp/envs/pgsr scripts/run_pgsr_benchmark.py" + + # --- 分支 1.2：表面法线与几何正则化 --- + "/root/autodl-tmp/env scripts/run_sugar_benchmark.py" + "/root/autodl-tmp/envs/erank_gs scripts/run_erankgs_benchmark.py" + + # --- 分支 1.3：隐式场与 SDF 融合 --- + "/root/autodl-tmp/envs/gof scripts/run_gof_benchmark.py" + "/root/autodl-tmp/envs/gspull scripts/run_gspull_benchmark.py" + + + # "/root/autodl-tmp/envs/cadgs scripts/run_qgs_benchmark.py" +) + +echo "🚀 SplatAtlas: Starting Robust Geometry Pipeline..." + +for task in "${tasks[@]}"; do + read -r env_path script_path <<< "$task" + + echo "----------------------------------------------------------------" + echo "🔄 Method: $(basename "$script_path")" + + # 1. 激活环境 + conda activate "$env_path" + + if [ $? -eq 0 ]; then + # 2. 关键修复：环境自动体检，确保 lpips 和 imageio 存在 + echo "🛠️ Checking requirements in $env_path..." + pip install lpips imageio -q + + # 3. 关键修复：强制使用当前环境的绝对路径 Python，防止 eval 泄露到 base 环境 + echo "✅ Environment ready. Executing..." + "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" + + if [ $? -eq 0 ]; then + echo "✨ Finished: $(basename "$script_path")" + else + echo "❌ Error: $(basename "$script_path") failed." + fi + else + echo "❌ Failed to activate: $env_path" + fi + + conda deactivate +done + +echo "================================================================" +echo "🏁 Geometry Pipeline completed." \ No newline at end of file diff --git a/scripts/run_all_benchmarks_2.sh b/scripts/run_all_benchmarks_2.sh new file mode 100644 index 0000000000000000000000000000000000000000..4a7e13ac0f10349ca8dfe4d8f1f599205ec33df0 --- /dev/null +++ b/scripts/run_all_benchmarks_2.sh @@ -0,0 +1,68 @@ +#!/bin/bash + +# 初始化 Conda 环境路径 +CONDA_BASE=$(conda info --base) +source "$CONDA_BASE/etc/profile.d/conda.sh" + +# 确保主路径在 PYTHONPATH 中 +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +# 定义任务队列：(环境绝对路径 脚本相对路径) +declare -a tasks=( + + "/root/autodl-tmp/envs/gaussianpro scripts/run_gaussianpro_benchmark.py" + + "/root/autodl-tmp/envs/conegs_env scripts/run_conegs_benchmark.py" + "/root/autodl-tmp/envs/steepgs_env scripts/run_steepgs.py" + + "/root/autodl-tmp/envs/absgs_env scripts/run_absgs_benchmark.py" + "/root/autodl-tmp/envs/reactgs_env scripts/run_reactgs_benchmark.py" + "/root/autodl-tmp/envs/lapis_gs scripts/run_lapisgs_benchmark.py" + "/root/autodl-tmp/envs/co_adaptation_3dgs scripts/run_coadaptgs_benchmark.py" + "/root/autodl-tmp/envs/mini_splatting scripts/run_minisplatting_benchmark.py" + "/root/autodl-tmp/envs/gs_lpm scripts/run_gslpm.py" + "/root/autodl-tmp/envs/opti3dgs scripts/run_opti3dgs_benchmark.py" + "/root/autodl-tmp/envs/ghap_env scripts/run_ghap_benchmark.py" + "/root/autodl-tmp/envs/3dgs_mcmc scripts/run_3dgsmcmc_benchmark.py" + # "/root/autodl-tmp/envs/lod_gs scripts/run_lod_gs_benchmark.py" + +#conegs可能需要重跑，test可能没有被划分 +) + +echo "🚀 SplatAtlas: Starting Robust Optimization Pipeline..." + +for task in "${tasks[@]}"; do + # 解析环境路径和脚本 + read -r env_path script_path <<< "$task" + + echo "----------------------------------------------------------------" + echo "🔄 Method: $(basename "$script_path")" + + # 1. 激活环境 + conda activate "$env_path" + + if [ $? -eq 0 ]; then + # 2. 环境自动体检：添加了 scikit-image 的安装 + echo "🛠️ Checking requirements in $env_path..." + # 增加 scikit-image 以解决 ModuleNotFoundError: No module named 'skimage' + pip install lpips imageio scikit-image -q + + # 3. 执行 Benchmark + echo "✅ Environment ready. Executing..." + "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" + + if [ $? -eq 0 ]; then + echo "✨ Finished: $(basename "$script_path")" + else + echo "❌ Error: $(basename "$script_path") failed." + fi + else + echo "❌ Failed to activate: $env_path" + fi + + # 释放当前环境 + conda deactivate +done + +echo "================================================================" +echo "🏁 Optimization Pipeline completed." \ No newline at end of file diff --git a/scripts/run_all_postprocessing.sh b/scripts/run_all_postprocessing.sh new file mode 100644 index 0000000000000000000000000000000000000000..178ee019a9e960035a39ca06b31a972ede4d11cb --- /dev/null +++ b/scripts/run_all_postprocessing.sh @@ -0,0 +1,18 @@ +#!/bin/bash +set -e + +cd "$(dirname "$0")/.." + +echo "🚀 [1/4] Extracting Dynamics..." +python scripts/extract_dynamics.py + +echo "🔬 [2/4] Computing Offline Physics..." +python scripts/compute_offline_physics.py + +echo "🌌 [3/4] Mapping Functional Space..." +python scripts/functional_clustering.py + +echo "📊 [4/4] Building Grand Leaderboard..." +python scripts/build_grand_table.py + +echo "✅ All post-processing pipelines completed successfully." diff --git a/scripts/run_analyticsplatting_benchmark.py b/scripts/run_analyticsplatting_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..ebe485a690fb72f1a6e0289355a2a524b9485836 --- /dev/null +++ b/scripts/run_analyticsplatting_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/analyticsplatting_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/analyticsplatting_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_atomgs_benchmark.py b/scripts/run_atomgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..375183eb303862a52c536db0853e9cf3f10019ba --- /dev/null +++ b/scripts/run_atomgs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/atomgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/AtomGS_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_bags_benchmark.py b/scripts/run_bags_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..022a8cabd53e2651a8bdbfeac52b7f7076f16778 --- /dev/null +++ b/scripts/run_bags_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" ❌ [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" ✨ [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" ❌ [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/bags_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/bags_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n🔥 [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" ✅ [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" ⚠️ [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" ✅ [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" ✅ [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"✅✅✅ {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_batch.py b/scripts/run_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..e357e3669cfe97aa6c942346a92efb3873151ced --- /dev/null +++ b/scripts/run_batch.py @@ -0,0 +1,88 @@ +import yaml +import os +import subprocess +import sys +import argparse +import concurrent.futures +import queue + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--config", type=str, default="configs/oracle_baseline.yaml") + parser.add_argument("--gpus", type=str, default="0") + parser.add_argument("--seeds", type=str, default="0", help="Comma separated seeds, e.g., 0,42,123") + return parser.parse_args() + +def run_task(task, config, gpu_queue, seed): + gpu_id = gpu_queue.get() + try: + method = task['method'] + scene = task['scene'] + source = task['source_path'] + prior_type = task.get('prior_type', 'none') + + global_iters = config['global_settings']['iterations'] + save_iters = config['global_settings']['save_iterations'] + save_iters_str = " ".join(map(str, save_iters)) + out_dir = f"outputs/{method}_{scene}_seed{seed}" + + final_ply = f"{out_dir}/point_cloud/iteration_{global_iters}/point_cloud.ply" + if not os.path.exists(final_ply): + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = str(gpu_id) + train_cmd = f"python scripts/main_train.py --method {method} --source_path {source} --model_path {out_dir} --seed {seed} --iterations {global_iters} --save_iterations {save_iters_str}" + subprocess.run(train_cmd, shell=True, check=True, env=env) + + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = str(gpu_id) + + for step in save_iters: + ply_path = f"{out_dir}/point_cloud/iteration_{step}/point_cloud.ply" + if not os.path.exists(ply_path): continue + + m_test = f"{out_dir}/metrics_test_iter{step}.json" + m_train = f"{out_dir}/metrics_train_iter{step}.json" + m_ablate = f"{out_dir}/metrics_ablation_iter{step}.json" + + if os.path.exists(m_test) and os.path.exists(m_train) and os.path.exists(m_ablate): + continue + + render_cmd = f"python scripts/main_render.py --method {method} --source_path {source} --model_path {out_dir} --iteration {step}" + subprocess.run(render_cmd, shell=True, check=True, env=env) + + is_final = (step == global_iters) + quarantine_flag = "" if is_final else "--skip_quarantine" + + # --- 核心改动：注入 --colmap_dir {source} --- + eval_test = f"python ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {out_dir}/renders_test_{step} --gt_dir {out_dir}/gt_test_{step} --ply_path {ply_path} --output_json {m_test} --prior_type {prior_type} --depth_dir {out_dir}/depths_test_{step} --colmap_dir {source} {quarantine_flag}" + subprocess.run(eval_test, shell=True, check=True, env=env) + + eval_train = f"python ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {out_dir}/renders_train_{step} --gt_dir {out_dir}/gt_train_{step} --ply_path {ply_path} --output_json {m_train} --prior_type {prior_type} --colmap_dir {source} --skip_quarantine" + subprocess.run(eval_train, shell=True, check=True, env=env) + + eval_ablate = f"python ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {out_dir}/renders_ablation_{step} --gt_dir {out_dir}/gt_test_{step} --ply_path {ply_path} --output_json {m_ablate} --prior_type {prior_type} --colmap_dir {source} --skip_quarantine" + subprocess.run(eval_ablate, shell=True, check=True, env=env) + + return True + except Exception: + return False + finally: + gpu_queue.put(gpu_id) + +def main(): + args = parse_args() + with open(args.config, 'r') as f: config = yaml.safe_load(f) + gpus = [g.strip() for g in args.gpus.split(',')] + gpu_queue = queue.Queue() + for g in gpus: gpu_queue.put(g) + + tasks = config.get('tasks', []) + seeds = [int(s.strip()) for s in args.seeds.split(",")] + with concurrent.futures.ThreadPoolExecutor(max_workers=len(gpus)) as executor: + futures = [executor.submit(run_task, task, config, gpu_queue, seed) for task in tasks for seed in seeds] + for future in concurrent.futures.as_completed(futures): pass + +if __name__ == "__main__": + main() diff --git a/scripts/run_bgtriangle_benchmark.py b/scripts/run_bgtriangle_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..2cff2a32843b5e38821ceace7c196631aaa68409 --- /dev/null +++ b/scripts/run_bgtriangle_benchmark.py @@ -0,0 +1,63 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/bgtriangle_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/bg_triangle_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_c3dgs_benchmark.py b/scripts/run_c3dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..895295b22133ff70d8f2731edaae33eff3918265 --- /dev/null +++ b/scripts/run_c3dgs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/c3dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/c3dgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"Processing Scene: {scene}") + + npz_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.npz") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(npz_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {npz_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_cdcgs_benchmark.py b/scripts/run_cdcgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..150d0d8f55f2c4e800eecdd4da3bc3396b3b4f83 --- /dev/null +++ b/scripts/run_cdcgs_benchmark.py @@ -0,0 +1,72 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Error] {name} Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [Success] {name} Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/cdcgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/CDC-GS_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_clone.sh b/scripts/run_clone.sh new file mode 100644 index 0000000000000000000000000000000000000000..968651ef6a911ad64187fa8be4a3c5166e7aa06e --- /dev/null +++ b/scripts/run_clone.sh @@ -0,0 +1,12 @@ +seconds=7200; \ +while [ $seconds -gt 0 ]; do \ + printf "\r倒计时: %02d:%02d:%02d" $((seconds/3600)) $((seconds%3600/60)) $((seconds%60)); \ + sleep 1; \ + : $((seconds--)); \ +done; \ +echo -e "\n开始运行..."; \ +while true; do \ + ./run_sgf_sequence.sh; \ + echo "脚本已结束，1分钟后重新运行..."; \ + sleep 60; \ +done \ No newline at end of file diff --git a/scripts/run_coadaptgs_benchmark.py b/scripts/run_coadaptgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..59a356a45aa9f8f9512eee8de07dad21095c1d37 --- /dev/null +++ b/scripts/run_coadaptgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/coadaptgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/CoAdaptGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_compact3dgs_benchmark.py b/scripts/run_compact3dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..89f8129a64efe210ad2bdede28e77d5313527091 --- /dev/null +++ b/scripts/run_compact3dgs_benchmark.py @@ -0,0 +1,84 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failure] {name} Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [Success] {name} Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/compact3dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Compact3DGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + npz_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.npz") + npz_pp_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud_pp.npz") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + has_checkpoint = os.path.exists(ply_path) or os.path.exists(npz_path) or os.path.exists(npz_pp_path) + + if has_checkpoint: + print(f" [Checkpoint] Checkpoint files found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"Complete: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_compgs.py b/scripts/run_compgs.py new file mode 100644 index 0000000000000000000000000000000000000000..0c725778594724534ba94c8ea949daf8fa88590d --- /dev/null +++ b/scripts/run_compgs.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Error] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [Success] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/compgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/CompGS_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_conegs_benchmark.py b/scripts/run_conegs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..6488ca151c9b3eae0d54725b3675f670e5bacbc3 --- /dev/null +++ b/scripts/run_conegs_benchmark.py @@ -0,0 +1,56 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Error] {name} Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [Success] {name} Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/conegs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/conegs:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} resolution={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + run_step("Training", train_cmd, curr_env) + + print(f"Scene {scene} execution routine finished.\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_contextgs_benchmark.py b/scripts/run_contextgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..ae8bb76b1552b2e3568e5c16949b347bb1e77da9 --- /dev/null +++ b/scripts/run_contextgs_benchmark.py @@ -0,0 +1,80 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/contextgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/ContextGS_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_drkgs_benchmark.py b/scripts/run_drkgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0a4833d7d7fef3254e52726e466ad18c1bf39723 --- /dev/null +++ b/scripts/run_drkgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/drkgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/DRK_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" Checkpoint Ply file found.") + if not os.path.exists(render_flag): + print(f" Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" Checkpoint Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" Checkpoint Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_eagles_benchmark.py b/scripts/run_eagles_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..8802560c7fbb31b5d1df2fb7a52c9885db09e9e2 --- /dev/null +++ b/scripts/run_eagles_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" FAILED: [{name}] Elapsed: {elapsed:.1f}s") + return False + print(f" SUCCESS: [{name}] Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" EXCEPTION: [{name}] {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/eagles_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/EAGLES_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n==== Dataset Domain: {ds['name']} ====") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} r={res}") + + pkl_path = os.path.join(model_out, f"iteration_{iters}/compressed.pkl") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(pkl_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {pkl_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_erankgs_benchmark.py b/scripts/run_erankgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..3f04fe9b4d577a67eeef70bb463d25a2181e1b11 --- /dev/null +++ b/scripts/run_erankgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/erankgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/erank_gs:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[{scene}] Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_eval_all_missing.py b/scripts/run_eval_all_missing.py new file mode 100644 index 0000000000000000000000000000000000000000..ad33e3c164bbacd730db449736485ee7ce25aade --- /dev/null +++ b/scripts/run_eval_all_missing.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +""" +Scan every directory under outputs/ that has a 30k ply but no metrics JSON, +infer (method, scene, resolution, conda_env) automatically, and run eval. + +Usage: + python scripts/run_eval_all_missing.py + python scripts/run_eval_all_missing.py --dry-run # print plan only +""" +import os +import sys +import glob +import argparse +import subprocess + +OUTPUTS_ROOT = "/root/autodl-tmp/SplatAtlas/outputs" +SPLATATLAS = "/root/autodl-tmp/SplatAtlas" +ITERS = 30000 + +# ── Scene → (dataset_path, resolution) ──────────────────────────────────────── +OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} +INDOOR_360 = {"bonsai", "counter", "kitchen", "room"} +TNT = { + "auditorium", "ballroom", "courtroom", "family", "francis", "horse", + "lighthouse", "m60", "museum", "palace", "panther", "playground", + "temple", "train" +} +ALL_SCENES = OUTDOOR_360 | INDOOR_360 | TNT + +def scene_info(scene): + if scene in OUTDOOR_360: + return f"/root/autodl-tmp/dataset/360/{scene}", 4 + if scene in INDOOR_360: + return f"/root/autodl-tmp/dataset/360/{scene}", 2 + return f"/root/autodl-tmp/dataset/tnt/{scene}", 2 + +# ── Method → conda env + PYTHONPATH ─────────────────────────────────────────── +def method_env(method): + if method.startswith("absgs"): + return "/root/autodl-tmp/envs/absgs_env", \ + f"{SPLATATLAS}:/root/autodl-tmp/absgs_official" + if method.startswith("pixelgs"): + return "/root/autodl-tmp/envs/pixelgs", \ + f"{SPLATATLAS}:/root/autodl-tmp/Pixel-GS" + # vanilla / vanillasgf / edgeloss / vanilla_thresh / etc. → base env + return "/root/miniconda3", \ + f"{SPLATATLAS}:/root/autodl-tmp/3dgsAtlas_official" + +# Only run eval for methods we own (skip pre-existing experiments like rdogaussian) +KNOWN_PREFIXES = ( + "vanillasgf", "vanilla_3dgs", "vanilla_logged", "vanilla_thresh", + "absgs", "pixelgs", "edgeloss", +) + +# ── Parse (method, scene) from directory name ────────────────────────────────── +def parse_dir(dirname): + """ + Handles patterns like: + vanillasgf_nosg_bicycle + vanilla_thresh_bicycle_t000028 + absgs_auditorium + Returns (method, scene) or (None, None). + """ + # Special case: vanilla_thresh has an extra threshold tag at the end + if dirname.startswith("vanilla_thresh_"): + rest = dirname[len("vanilla_thresh_"):] # e.g. "bicycle_t000028" + for scene in ALL_SCENES: + if rest == scene or rest.startswith(scene + "_"): + return "vanilla_thresh", scene + return None, None + + # General: strip tokens from the right until we find a known scene + tokens = dirname.split("_") + for i in range(len(tokens) - 1, 0, -1): + candidate_scene = "_".join(tokens[i:]) # handles multi-word scene names + candidate_method = "_".join(tokens[:i]) + if candidate_scene in ALL_SCENES: + return candidate_method, candidate_scene + + return None, None + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--dry-run", action="store_true", + help="Print what would be run without executing") + args = ap.parse_args() + + dirs = sorted(os.listdir(OUTPUTS_ROOT)) + to_eval = [] + + for dirname in dirs: + dirpath = os.path.join(OUTPUTS_ROOT, dirname) + if not os.path.isdir(dirpath): + continue + + ply = os.path.join(dirpath, f"point_cloud/iteration_{ITERS}/point_cloud.ply") + metrics = os.path.join(dirpath, f"metrics_test_iter{ITERS}.json") + + if not os.path.exists(ply): + continue # training not done + if os.path.exists(metrics): + continue # eval already done + + method, scene = parse_dir(dirname) + + if method is None or scene is None: + print(f"SKIP (unparseable) : {dirname}") + continue + + if not any(method.startswith(p) for p in KNOWN_PREFIXES): + print(f"SKIP (not ours) : {dirname} method={method}") + continue + + to_eval.append((dirname, method, scene, dirpath)) + + print(f"\nNeed eval: {len(to_eval)} directories\n") + if not to_eval: + print("Nothing to do.") + return + + failed = [] + for dirname, method, scene, dirpath in to_eval: + ply = os.path.join(dirpath, f"point_cloud/iteration_{ITERS}/point_cloud.ply") + metrics = os.path.join(dirpath, f"metrics_test_iter{ITERS}.json") + renders = os.path.join(dirpath, f"renders_test_{ITERS}") + gt_dir = os.path.join(dirpath, f"gt_test_{ITERS}") + depths = os.path.join(dirpath, f"depths_test_{ITERS}") + + source, res = scene_info(scene) + conda_env, pp = method_env(method) + + cmd = ( + f"conda run --no-capture-output -p {conda_env} " + f"python {SPLATATLAS}/ufd_evalkit/run_eval.py " + f"--method {method} --scene {scene} " + f"--render_dir {renders} " + f"--gt_dir {gt_dir} " + f"--ply_path {ply} " + f"--output_json {metrics} " + f"--colmap_dir {source} " + f"--depth_dir {depths} " + ) + + print(f"{'[DRY]' if args.dry_run else '===>'} {dirname} (method={method}, scene={scene}, res={res})") + if args.dry_run: + continue + + env = os.environ.copy() + env["PYTHONPATH"] = pp + result = subprocess.run(cmd, shell=True, env=env) + + if result.returncode == 0: + print(f" OK") + else: + print(f" FAILED rc={result.returncode}") + failed.append(dirname) + + if not args.dry_run: + print(f"\n{'='*60}") + print(f"Done. Failed: {len(failed)}") + for d in failed: + print(f" FAIL {d}") + if failed: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/scripts/run_fds_benchmark.py b/scripts/run_fds_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0e6e84b45661a5532b3ee001e11675932bf69aef --- /dev/null +++ b/scripts/run_fds_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/fds_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/FDS-GS_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n Dataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Table Update", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_fdsgs_benchmark.py b/scripts/run_fdsgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0f89160711266222a631393b812921cac9b9df09 --- /dev/null +++ b/scripts/run_fdsgs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/fdsgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/FDS-GS_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n==== Dataset: {ds['name']} ====") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n Processing Scene: {scene} (res={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Update Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_flod_benchmark.py b/scripts/run_flod_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..6ff2e96ad670c9eb336a277c7fa343d6b42725e2 --- /dev/null +++ b/scripts/run_flod_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/flod_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/flod_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" Ply file found.") + if not os.path.exists(render_flag): + print(f" Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gaussianfocus.py b/scripts/run_gaussianfocus.py new file mode 100644 index 0000000000000000000000000000000000000000..8ea97ca4b8a25faddccb79381202b8149a1abbd1 --- /dev/null +++ b/scripts/run_gaussianfocus.py @@ -0,0 +1,370 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [STEP] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [FAILED] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gaussianfocus_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GaussianFocus_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[INFO] [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [INFO] [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [CLEANUP] [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [INFO] [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [INFO] [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[DONE] {scene} Done!\n") + +if __name__ == "__main__": + main() +EOFcat << 'EOF' > /root/autodl-tmp/SplatAtlas/methods/wrapper_gaussianfocus.py +import os +import sys +import random +import torch +import numpy as np +import torch.nn.functional as F +import torch.nn as nn +import torch.optim as optim +from argparse import ArgumentParser + +from core.registry import register_method +from core.base_method import BaseMethod + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../GaussianFocus_offy'))) +from utils.loss_utils import l1_loss, ssim +from gaussian_renderer import render as native_render +from scene import Scene, GaussianModel +from arguments import ModelParams, PipelineParams, OptimizationParams + +def edge_loss(input_t, target_t): + edge_filter = torch.Tensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]]).view(1, 1, 3, 3).to(input_t.device) + edge_filter = edge_filter.repeat(input_t.size(1), 1, 1, 1) + edge_filter = nn.Parameter(data=edge_filter, requires_grad=False) + input_x = F.conv2d(input_t, edge_filter, padding=1, groups=input_t.size(1)) + input_y = F.conv2d(input_t, edge_filter.transpose(2, 3), padding=1, groups=input_t.size(1)) + target_x = F.conv2d(target_t, edge_filter, padding=1, groups=target_t.size(1)) + target_y = F.conv2d(target_t, edge_filter.transpose(2, 3), padding=1, groups=target_t.size(1)) + loss_x = F.l1_loss(input_x, target_x) + loss_y = F.l1_loss(input_y, target_y) + return (loss_x + loss_y) / 2.0 + +def frequency_loss(input_t, target_t): + grad_x_input = input_t[:, :, 1:] - input_t[:, :, :-1] + grad_y_input = input_t[:, 1:, :] - input_t[:, :-1, :] + grad_x_target = target_t[:, :, 1:] - target_t[:, :, :-1] + grad_y_target = target_t[:, 1:, :] - target_t[:, :-1, :] + loss_x = F.l1_loss(grad_x_input, grad_x_target) + loss_y = F.l1_loss(grad_y_input, grad_y_target) + return (loss_x + loss_y) / 2.0 + +class CombinedLoss(nn.Module): + def __init__(self): + super(CombinedLoss, self).__init__() + + def forward(self, finals, gt_image, beta=1.0, eta=1.0): + if len(finals.shape) == 3: + finals = finals.unsqueeze(0) + if len(gt_image.shape) == 3: + gt_image = gt_image.unsqueeze(0) + e_loss = edge_loss(finals, gt_image) + f_loss = frequency_loss(finals, gt_image) + return beta * e_loss + eta * f_loss + +class SimplePatchattn(nn.Module): + def __init__(self): + super(SimplePatchattn, self).__init__() + self.conv1 = nn.Conv2d(6, 16, kernel_size=3, padding=1) + self.conv2 = nn.Conv2d(16, 16, kernel_size=3, padding=1) + self.conv3 = nn.Conv2d(16, 3, kernel_size=3, padding=1) + + def forward(self, image1, image2, block_size=64): + _, _, H, W = image1.shape + blocks = [] + for i in range(0, H, block_size): + for j in range(0, W, block_size): + block1 = image1[:, :, i:i + block_size, j:j + block_size] + block2 = image2[:, :, i:i + block_size, j:j + block_size] + pad_height = block_size - block1.size(2) + pad_width = block_size - block1.size(3) + if pad_height > 0 or pad_width > 0: + block1 = F.pad(block1, (0, pad_width, 0, pad_height)) + block2 = F.pad(block2, (0, pad_width, 0, pad_height)) + block = torch.cat((block1, block2), dim=1) + block = F.relu(self.conv1(block)) + block = F.relu(self.conv2(block)) + block = self.conv3(block) + blocks.append(block[:, :, :min(block_size, H - i), :min(block_size, W - j)]) + num_blocks_per_row = W // block_size + (W % block_size > 0) + rows = [torch.cat(blocks[i * num_blocks_per_row:(i + 1) * num_blocks_per_row], dim=3) for i in range(H // block_size + (H % block_size > 0))] + return torch.cat(rows, dim=2) + +@register_method("gaussianfocus") +class GaussianFocusWrapper(BaseMethod): + def __init__(self, dataset_config, hyperparams): + self.parser = ArgumentParser() + self.lp = ModelParams(self.parser) + self.op = OptimizationParams(self.parser) + self.pp = PipelineParams(self.parser) + self.args = self.parser.parse_args([]) + self.args.source_path = dataset_config["source_path"] + self.args.model_path = dataset_config["model_path"] + self.args.eval = True + self.args.resolution = dataset_config.get("resolution", 1) + self.track_decoupling = hyperparams.get("track_decoupling", False) + self.dataset = self.lp.extract(self.args) + self.opt = self.op.extract(self.args) + self.pipe = self.pp.extract(self.args) + self.gaussians = GaussianModel(self.dataset.sh_degree) + self.scene = Scene(self.dataset, self.gaussians) + self.gaussians.training_setup(self.opt) + bg_color = [1, 1, 1] if self.dataset.white_background else [0, 0, 0] + self.background = torch.tensor(bg_color, dtype=torch.float32, device="cuda") + self.viewpoint_stack = self.scene.getTrainCameras().copy() + self.last_n_gaussians = len(self.gaussians.get_xyz) + self.patch_attention = SimplePatchattn().to("cuda") + self.patch_optimizer = optim.Adam(self.patch_attention.parameters(), lr=1e-3) + self.combined_loss = CombinedLoss().to("cuda") + + def train_iteration(self, step): + self.gaussians.update_learning_rate(step) + if step % 1000 == 0: + self.gaussians.oneupSHdegree() + if not self.viewpoint_stack: + self.viewpoint_stack = self.scene.getTrainCameras().copy() + viewpoint_cam = self.viewpoint_stack.pop(random.randint(0, len(self.viewpoint_stack) - 1)) + subpixel_offset = None + if hasattr(self.dataset, "ray_jitter") and self.dataset.ray_jitter: + subpixel_offset = torch.rand((int(viewpoint_cam.image_height), int(viewpoint_cam.image_width), 2), dtype=torch.float32, device="cuda") - 0.5 + render_pkg = native_render(viewpoint_cam, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size, subpixel_offset=subpixel_offset) + image = render_pkg["render"] + viewspace_point_tensor = render_pkg["viewspace_points"] + visibility_filter = render_pkg["visibility_filter"] + radii = render_pkg["radii"] + gt_image = viewpoint_cam.original_image.cuda() + if step % 50 == 0: + image1 = image.unsqueeze(0) + gt_image1 = gt_image.unsqueeze(0) + output = self.patch_attention(image1, gt_image1) + finals = output * image + torch.cuda.empty_cache() + else: + finals = image + finals_sq = torch.squeeze(finals, 0) if finals.dim() == 4 else finals + Ll1 = l1_loss(image, gt_image) + ssim_value = ssim(image, gt_image) + loss_target = (1.0 - self.opt.lambda_dssim) * Ll1 + self.opt.lambda_dssim * (1.0 - ssim_value) + loss_parasitic = self.combined_loss(finals_sq, gt_image) + loss = loss_target + loss_parasitic + grad_cos_sim = 0.0 + parasitic_ratio = 0.0 + metrics_sti = {} + if self.track_decoupling and step % 100 == 0: + param_groups_map = { + "spatial": [self.gaussians._xyz], + "geometry": [self.gaussians._scaling, self.gaussians._rotation], + "opacity": [self.gaussians._opacity], + "appearance": [self.gaussians._features_dc, self.gaussians._features_rest], + } + def get_effective_step(loss_val): + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss_val.backward(retain_graph=True) + step_dict = {} + for g_name, p_list in param_groups_map.items(): + u_list = [] + for p in p_list: + if p.grad is not None: + g = p.grad.clone() + state = self.gaussians.optimizer.state.get(p, None) + v = state['exp_avg_sq'] if (state is not None and 'exp_avg_sq' in state) else torch.ones_like(g) * 1e-8 + lr = 1e-3 + for opt_g in self.gaussians.optimizer.param_groups: + if any(id(opt_p) == id(p) for opt_p in opt_g['params']): + lr = opt_g['lr'] + break + u = (lr / (torch.sqrt(v) + 1e-15)) * g + u_list.append(u.reshape(-1)) + if u_list: + step_dict[g_name] = torch.cat(u_list) + return step_dict + steps_target = get_effective_step(loss_target) + steps_parasitic = get_effective_step(loss_parasitic) + u_t_global = [] + u_p_global = [] + for g_name in param_groups_map: + ut = steps_target.get(g_name) + up = steps_parasitic.get(g_name) + if ut is not None and up is not None and ut.norm() > 0 and up.norm() > 0: + u_t_global.append(ut) + u_p_global.append(up) + cos_val = float(F.cosine_similarity(ut.unsqueeze(0), up.unsqueeze(0))) + r_val = float(up.norm() / (ut.norm() + up.norm() + 1e-7)) + ti = r_val * max(0.0, -cos_val) + metrics_sti[f"sti_{g_name}"] = ti + else: + metrics_sti[f"sti_{g_name}"] = 0.0 + if u_t_global and u_p_global: + ut_cat = torch.cat(u_t_global) + up_cat = torch.cat(u_p_global) + grad_cos_sim = float(F.cosine_similarity(ut_cat.unsqueeze(0), up_cat.unsqueeze(0))) + parasitic_ratio = float(up_cat.norm() / (ut_cat.norm() + 1e-7)) + if step % 50 == 0: + self.patch_optimizer.zero_grad() + self.gaussians.optimizer.zero_grad(set_to_none=True) + loss.backward() + else: + if step % 50 == 0: + self.patch_optimizer.zero_grad() + loss.backward() + with torch.no_grad(): + if step < self.opt.densify_until_iter: + self.gaussians.max_radii2D[visibility_filter] = torch.max(self.gaussians.max_radii2D[visibility_filter], radii[visibility_filter]) + self.gaussians.add_densification_stats(viewspace_point_tensor, visibility_filter) + if step > self.opt.densify_from_iter and step % self.opt.densification_interval == 0: + size_threshold = 0.02 if step > 1000 else None + self.gaussians.densify_and_prune(self.opt.densify_grad_threshold, 0.005, self.scene.cameras_extent, size_threshold, step) + if step % self.opt.opacity_reset_interval == 0 or (self.dataset.white_background and step == self.opt.densify_from_iter): + self.gaussians.reset_opacity() + self.gaussians.optimizer.step() + self.gaussians.optimizer.zero_grad(set_to_none=True) + if step % 50 == 0: + self.patch_optimizer.step() + self.patch_optimizer.zero_grad() + num_gaussians = self.gaussians.get_xyz.shape[0] + metrics = { + "loss": float(loss), "loss_l1": float(loss_target), "loss_ssim": float(loss_parasitic), + "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians), + "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)), + "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio) + } + metrics.update(metrics_sti) + self.last_n_gaussians = num_gaussians + histograms = {} + if step % 1000 == 0: + histograms["opacity"] = torch.sigmoid(self.gaussians._opacity).clone().detach() + scales = torch.exp(self.gaussians._scaling).clone().detach() + histograms["scaling"] = scales + gamma = scales.max(dim=-1)[0] / (scales.min(dim=-1)[0] + 1e-7) + histograms["anisotropy"] = gamma + histograms["sh_dc_mag"] = self.gaussians._features_dc.detach().norm(dim=-1) + return metrics, histograms + + def render(self, camera): + with torch.no_grad(): + subpixel_offset = None + render_pkg = native_render(camera, self.gaussians, self.pipe, self.background, kernel_size=self.dataset.kernel_size, subpixel_offset=subpixel_offset) + return {"image": render_pkg["render"], "depth": render_pkg.get("depth", None)} + + def save(self, save_dir, step): + self.scene.save(step) + + def load(self, model_path, iteration): + self.gaussians.load_ply(os.path.join(model_path, 'point_cloud', f'iteration_{iteration}', 'point_cloud.ply')) + + def get_spatial_centers(self): + return self.gaussians._xyz + + def compute_physical_metrics(self, cameras=None): + metrics = {} + with torch.no_grad(): + raw_scales = self.gaussians._scaling + scales = torch.exp(raw_scales) + max_S, _ = torch.max(scales, dim=1) + min_S, _ = torch.min(scales, dim=1) + gamma = max_S / (min_S + 1e-7) + metrics["gamma_median"] = float(torch.median(gamma)) + metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90)) + metrics["scale_mean"] = float(torch.mean(scales)) + metrics["alpha_mean"] = float(torch.mean(torch.sigmoid(self.gaussians._opacity))) + dc, rest = self.gaussians._features_dc, self.gaussians._features_rest + if rest is not None and rest.shape[1] > 0: + metrics["sh_energy_ratio"] = float(rest.norm(dim=-1).mean() / (dc.norm(dim=-1).mean() + 1e-7)) + if cameras is not None and len(cameras) > 0: + view_dirs = [] + for c in cameras: + view_dirs.append(c.world_view_transform[:3, 2].tolist()) + view_dirs = F.normalize(torch.tensor(view_dirs, dtype=torch.float32, device="cuda"), dim=1) + rots = F.normalize(self.gaussians._rotation.clone(), dim=1) + w, x, y, z = rots.unbind(dim=-1) + normals = F.normalize(torch.stack([2*(x*z + w*y), 2*(y*z - w*x), 1-2*(x*x + y*y)], dim=-1), dim=1) + max_cos, _ = torch.max(torch.abs(torch.matmul(normals, view_dirs.T)), dim=1) + metrics["billboard_bias_ratio"] = float((max_cos > 0.90).float().mean()) + return metrics + + def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor: + with torch.no_grad(): + V = query_points.shape[0] + densities = torch.zeros(V, device="cuda") + xyz, opacities = self.gaussians._xyz, torch.sigmoid(self.gaussians._opacity).squeeze() + scales = torch.exp(self.gaussians._scaling) + sigma_sq = scales.max(dim=1)[0].pow(2) + N_gaussians = xyz.shape[0] + chunk_size = max(1, 30_000_000 // (N_gaussians + 1)) + for i in range(0, V, chunk_size): + end = min(i + chunk_size, V) + dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2) + weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7)) + densities[i:end] = torch.sum(weights * opacities.unsqueeze(0), dim=1) + return densities diff --git a/scripts/run_gaussianfocus_benchmark.py b/scripts/run_gaussianfocus_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..2516b77288cf40fe0b63ac99776b366b5b2b37ed --- /dev/null +++ b/scripts/run_gaussianfocus_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Fail] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [Pass] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Fail] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gaussianfocus_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GaussianFocus_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[Processing] Method: {method.upper()} | Scene: {scene} | Res: {res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[Done] {scene} finished.\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gaussianpro_benchmark.py b/scripts/run_gaussianpro_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..fe6d05dd1f33970009ee8b6e502ed08903c28b4e --- /dev/null +++ b/scripts/run_gaussianpro_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name} {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gaussianpro_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gaussianpro_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"Done scene: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gaussianshader.py b/scripts/run_gaussianshader.py new file mode 100644 index 0000000000000000000000000000000000000000..57ef37a216bb5f72372381f40152c37a58d9e326 --- /dev/null +++ b/scripts/run_gaussianshader.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gaussianshader_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gaussian_shader_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gaussianshader_benchmark.py b/scripts/run_gaussianshader_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..c25bb5f1b287638988b4dffc08b405a8bbb97a19 --- /dev/null +++ b/scripts/run_gaussianshader_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gaussianshader_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GaussianShader:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gaussianwild_benchmark.py b/scripts/run_gaussianwild_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..b76ee944c344fb70a2f54a824b517d6a11724d6c --- /dev/null +++ b/scripts/run_gaussianwild_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f"[Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f"[ERROR] {name} Failed. Elapsed: {elapsed:.1f}s") + return False + print(f"[SUCCESS] {name} Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f"[EXCEPTION] {name} Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gaussianwild_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Gaussian-Wild_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print("[INFO] Checkpoint Ply file found.") + if not os.path.exists(render_flag): + print("[INFO] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print("[INFO] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print("[INFO] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[SUCCESS] {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_geogaussian_benchmark.py b/scripts/run_geogaussian_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..b4e77c134438c8e96986d92395be34233e311ab0 --- /dev/null +++ b/scripts/run_geogaussian_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/geogaussian_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GeoGaussian:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" Ply file found.") + if not os.path.exists(render_flag): + print(f" Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_ges_benchmark.py b/scripts/run_ges_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..cc7184b7116c6e46fab8c65dc7d346fead07831d --- /dev/null +++ b/scripts/run_ges_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [COMPLETED] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/ges_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/ges_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"DONE: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_ghap_benchmark.py b/scripts/run_ghap_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..9ee034a017bc03b5f59089de5d1aa0ac5a396dd0 --- /dev/null +++ b/scripts/run_ghap_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/ghap_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/ghap_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gigs_benchmark.py b/scripts/run_gigs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..527e603b4949daef28a617e91c27d2a169f069e6 --- /dev/null +++ b/scripts/run_gigs_benchmark.py @@ -0,0 +1,82 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gigs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GI-GS_offy:{curr_env.get('PYTHONPATH', '')}" + curr_env["TORCH_CUDA_ARCH_LIST"] = "8.6" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gof_benchmark.py b/scripts/run_gof_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..2f01277973cfeeb9f04a00a5bd2342c5ad797be0 --- /dev/null +++ b/scripts/run_gof_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" FAILED: [{name}] Elapsed: {elapsed:.1f}s") + return False + print(f" COMPLETED: [{name}] Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" EXCEPTION: [{name}] {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gof_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gof_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDOMAIN: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nPROCESSING: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" FOUND: Checkpoint") + if not os.path.exists(render_flag): + print(" CLEANUP: Re-rendering") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" FOUND: Render assets") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" FOUND: Metrics JSON") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Updating Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"DONE: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gsdf_benchmark.py b/scripts/run_gsdf_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..6a192e49994912cea7f4bd10cd1eeff394d5747e --- /dev/null +++ b/scripts/run_gsdf_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gsdf_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GSDF_run:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gsgn_benchmark.py b/scripts/run_gsgn_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..946fcdbf5d1767382567a366130fd3d7ff5bdf2e --- /dev/null +++ b/scripts/run_gsgn_benchmark.py @@ -0,0 +1,74 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gsgn_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gsgn_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n==== Dataset Domain: {ds['name']} ====") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[{scene}] Processing Finished.\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gslpm.py b/scripts/run_gslpm.py new file mode 100644 index 0000000000000000000000000000000000000000..40c3c1d3925394bf8fd92673ab8cedfe0607e859 --- /dev/null +++ b/scripts/run_gslpm.py @@ -0,0 +1,84 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f"[{name}] Starting...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f"[{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f"[{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f"[{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gslpm_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GS-LPM:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print("Checkpoint found.") + if not os.path.exists(render_flag): + print("Re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): + shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("OfflineRendering", render_cmd, curr_env): + continue + else: + print("Render complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): + continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics", eval_cmd, curr_env) + else: + print("Metrics JSON exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("OfflinePhysics", physics_cmd, curr_env) + + run_step("GrandTable", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done.") + +if __name__ == "__main__": + main() diff --git a/scripts/run_gspull_benchmark.py b/scripts/run_gspull_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..a08a64a5e2c8bd1c0cba2ddf1a39bcba5f3db1fe --- /dev/null +++ b/scripts/run_gspull_benchmark.py @@ -0,0 +1,53 @@ +import os +import yaml +import subprocess +import time + +def run_step(name, cmd, env): + print(f"\n [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gspull_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = "/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/GS-Pull:/root/autodl-tmp/GS-Pull/gaussian_splatting" + curr_env["LD_LIBRARY_PATH"] = "/root/autodl-tmp/envs/gspull/lib:" + curr_env.get("LD_LIBRARY_PATH", "") + curr_env["CUDA_VISIBLE_DEVICES"] = curr_env.get("CUDA_VISIBLE_DEVICES", "0") + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + train_cmd = f"/root/autodl-tmp/envs/gspull/bin/python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + physics_cmd = f"/root/autodl-tmp/envs/gspull/bin/python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Physics", physics_cmd, curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_gsshader_benchmark.py b/scripts/run_gsshader_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0eb5338e0e1a691544fc2feecfc48797f3a9abee --- /dev/null +++ b/scripts/run_gsshader_benchmark.py @@ -0,0 +1,88 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + # 增加 check=True 捕获更底层的异常 + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" ❌ [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" ✨ [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" ❌ [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/gsshader_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + # 注入环境补丁以解决 MKL 冲突 + curr_env = os.environ.copy() + curr_env["MKL_SERVICE_FORCE_INTEL"] = "1" + curr_env["MKL_THREADING_LAYER"] = "GNU" + # 显式指向当前虚拟环境的 lib 路径 + conda_lib_path = "/root/autodl-tmp/envs/gaussianshader/lib" + curr_env["LD_LIBRARY_PATH"] = f"{conda_lib_path}:{curr_env.get('LD_LIBRARY_PATH', '')}" + + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gsshader_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n🔥 [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" ✅ [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" ⚠️ [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"/root/autodl-tmp/envs/gaussianshader/bin/python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" ✅ [Checkpoint] Render assets complete.") + else: + train_cmd = f"/root/autodl-tmp/envs/gaussianshader/bin/python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"/root/autodl-tmp/envs/gaussianshader/bin/python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" ✅ [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"/root/autodl-tmp/envs/gaussianshader/bin/python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"✅✅✅ {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_hac_plus_benchmark.py b/scripts/run_hac_plus_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..5d5947fc7b731a63b9cefab84d1d13ddc00fa0e6 --- /dev/null +++ b/scripts/run_hac_plus_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [COMPLETED] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/hac_plus_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/hac_plus_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_hogs_benchmark.py b/scripts/run_hogs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..c8088524c3221197dfa41cb081609ed7b6d97b02 --- /dev/null +++ b/scripts/run_hogs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/hogs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/HoGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_improvingadc_benchmark.py b/scripts/run_improvingadc_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0062931f27277f7f6bde55fcc56370c460c06ec5 --- /dev/null +++ b/scripts/run_improvingadc_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/improvingadc_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/improvingadc_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_lapisgs_benchmark.py b/scripts/run_lapisgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..26b5ead2ad0bc2199027a0de3da77a0298431adc --- /dev/null +++ b/scripts/run_lapisgs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/lapisgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/LapisGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_lightgaussian_benchmark.py b/scripts/run_lightgaussian_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..7f2320d9c7a91e6080534a5648a2fc0fc2c21c08 --- /dev/null +++ b/scripts/run_lightgaussian_benchmark.py @@ -0,0 +1,66 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + return False + return True + except Exception: + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/lightgaussian_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/LightGaussian_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Physics", physics_cmd, curr_env) + + run_step("Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_lod_gs_benchmark.py b/scripts/run_lod_gs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..1f8d7b70879d03eac6b05502b3ba1c78c6825f4c --- /dev/null +++ b/scripts/run_lod_gs_benchmark.py @@ -0,0 +1,83 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/lod_gs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + cap = cfg["global_settings"].get("cap_gaussians", None) + cap_flag = f" --cap_gaussians {cap}" if cap else "" + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/LOD_GS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling{cap_flag}" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_lp3dgs_benchmark.py b/scripts/run_lp3dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..bb3ad7ed34f4bb82bff2ef86b50c85ed11f58153 --- /dev/null +++ b/scripts/run_lp3dgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [COMPLETED] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/lp3dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/LP-3DGS_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[{scene}] Process Done.\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_minisplatting_benchmark.py b/scripts/run_minisplatting_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..df07289427d19d68bc0d07913488499660d74ed7 --- /dev/null +++ b/scripts/run_minisplatting_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f"[Step] {name}") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f"[Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f"[Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f"[Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/minisplatting_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/mini_splatting_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"Dataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print("Checkpoint Ply file found.") + if not os.path.exists(render_flag): + print("Cleanup Previous render incomplete, re-rendering.") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print("Checkpoint Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print("Checkpoint Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done.") + +if __name__ == "__main__": + main() diff --git a/scripts/run_minisplattingsgf.py b/scripts/run_minisplattingsgf.py new file mode 100644 index 0000000000000000000000000000000000000000..0eb46a033a4d1b2fab29c3eb4b2b9edb37a5d1d1 --- /dev/null +++ b/scripts/run_minisplattingsgf.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f"[Step] {name}") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f"[Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f"[Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f"[Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/minisplattingsgf_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/minisplatting_sgf_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"Dataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print("Checkpoint Ply file found.") + if not os.path.exists(render_flag): + print("Cleanup Previous render incomplete, re-rendering.") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print("Checkpoint Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print("Checkpoint Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done.") + +if __name__ == "__main__": + main() diff --git a/scripts/run_mipsplatting_benchmark.py b/scripts/run_mipsplatting_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..e9c486dfc14819f53e718d8e6a9f6c0fe2caa8b5 --- /dev/null +++ b/scripts/run_mipsplatting_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/mipsplatting_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/mip-splatting:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_mipsplattingsgf.py b/scripts/run_mipsplattingsgf.py new file mode 100644 index 0000000000000000000000000000000000000000..98a947b81d23153baef1b2176a0eeafca51ba9fc --- /dev/null +++ b/scripts/run_mipsplattingsgf.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/mipsplattingsgf_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/mip-splatting:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_new_methods-Copy1.sh b/scripts/run_new_methods-Copy1.sh new file mode 100644 index 0000000000000000000000000000000000000000..df512a28287f643b253191bd35ba446809e62130 --- /dev/null +++ b/scripts/run_new_methods-Copy1.sh @@ -0,0 +1,86 @@ +#!/bin/bash + +source /root/miniconda3/etc/profile.d/conda.sh + +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +ERROR_LOG="/root/autodl-tmp/SplatAtlas/run_errors.log" + +echo "================================================================" >> "$ERROR_LOG" +echo "New Run Started at $(date)" >> "$ERROR_LOG" +echo "================================================================" >> "$ERROR_LOG" + +declare -a tasks=( +# "/root/autodl-tmp/envs/lapis_gs scripts/run_lapisgs_benchmark.py" +# "/root/autodl-tmp/envs/mip-splatting scripts/run_mipsplatting_benchmark.py" +# --- Apr 13 --- +# "/root/autodl-tmp/envs/c3dgs scripts/run_c3dgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/compgs scripts/run_compgs.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/sogs scripts/run_sogs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/compact3dgs scripts/run_compact3dgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/context_gs scripts/run_contextgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/spec_gaussian scripts/run_specgaussian_benchmark.py" # 重，跳过 +# "/root/autodl-tmp/envs/fds_gs scripts/run_fdsgs_benchmark.py" # 已OOM，跳过 +# "/root/autodl-tmp/envs/cdc_gs scripts/run_cdcgs_benchmark.py" +"/root/autodl-tmp/envs/ngs scripts/run_ngs_benchmark.py" +# "/root/autodl-tmp/envs/aaa_gaussians scripts/run_aaags_benchmark.py" +# "/root/autodl-tmp/envs/sort_free_gs scripts/run_sortfreegs_benchmark.py" +# "/root/autodl-tmp/envs/analytic_splatting scripts/run_analyticsplatting_benchmark.py" +# "/root/autodl-tmp/envs/rade_gs scripts/run_radegs_benchmark.py" +# "/root/autodl-tmp/envs/splatfacto_w scripts/run_splatfactow_benchmark.py" # nerfstudio体系，结构不同，单独跑 +# "/root/autodl-tmp/envs/sundae scripts/run_sundae_benchmark.py" # 重，跳过 +# "/root/autodl-tmp/envs/pup_3dgs scripts/run_pup3dgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/improving_adc scripts/run_improvingadc_benchmark.py" +# "/root/autodl-tmp/envs/topology_gs scripts/run_topologygs_benchmark.py" +# "/root/autodl-tmp/envs/GeoGaussian scripts/run_geogaussian_benchmark.py" # 重，跳过 +# "/root/autodl-tmp/envs/FLoD scripts/run_flod_benchmark.py" # 多LOD重，跳过 +# "/root/autodl-tmp/envs/VDGS scripts/run_vdgs_benchmark.py" +# # # --- Apr 14 --- +# # "/root/autodl-tmp/envs/gspull scripts/run_gspull_benchmark.py" +# "/root/autodl-tmp/envs/gaussianshader scripts/run_gsshader_benchmark.py" +# "/root/autodl-tmp/envs/gaussian_surfels scripts/run_surfel_benchmark.py" # 重，跳过 +"/root/autodl-tmp/envs/gaussianfocus scripts/run_gaussianfocus_benchmark.py" +"/root/autodl-tmp/envs/gigs scripts/run_gigs_benchmark.py" +"/root/autodl-tmp/envs/bags scripts/run_bags_benchmark.py" +"/root/autodl-tmp/envs/drk scripts/run_drkgs_benchmark.py" +# "/root/autodl-tmp/envs/fdsgs scripts/run_fds_benchmark.py" # 已OOM +"/root/autodl-tmp/envs/op43dgs scripts/run_op43dgs_benchmark.py" +"/root/autodl-tmp/envs/lp_3dgs scripts/run_lp3dgs_benchmark.py" +) + +echo "SplatAtlas: Starting New Methods Pipeline..." +echo "Errors will be logged to: $ERROR_LOG" + +for task in "${tasks[@]}"; do + read -r env_path script_path <<< "$task" + method_name=$(basename "$script_path") + + echo "----------------------------------------------------------------" + echo "Method: $method_name" + + conda activate "$env_path" + + if [ $? -eq 0 ]; then + echo "Checking requirements in $env_path..." + pip install lpips imageio scikit-image -q + + echo "Environment ready. Executing..." + + PYTHONPATH=/root/autodl-tmp/SplatAtlas "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" 2> >(tee -a "$ERROR_LOG" >&2) + + if [ $? -eq 0 ]; then + echo "Finished: $method_name" + else + echo "Error: $method_name failed. Saved to $ERROR_LOG" + echo "[$(date)] FAILED: $method_name" >> "$ERROR_LOG" + fi + else + echo "Failed to activate: $env_path" + echo "[$(date)] FAILED (Env Activation): $env_path for $method_name" >> "$ERROR_LOG" + fi + + conda deactivate +done + +echo "================================================================" +echo "Pipeline completed. Check $ERROR_LOG for crash reports." diff --git a/scripts/run_new_methods.sh b/scripts/run_new_methods.sh new file mode 100644 index 0000000000000000000000000000000000000000..e4ed58615b63ad81d1d1e28a8c3061e99f038c3a --- /dev/null +++ b/scripts/run_new_methods.sh @@ -0,0 +1,86 @@ +#!/bin/bash + +source /root/miniconda3/etc/profile.d/conda.sh + +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +ERROR_LOG="/root/autodl-tmp/SplatAtlas/run_errors.log" + +echo "================================================================" >> "$ERROR_LOG" +echo "New Run Started at $(date)" >> "$ERROR_LOG" +echo "================================================================" >> "$ERROR_LOG" + +declare -a tasks=( +# "/root/autodl-tmp/envs/lapis_gs scripts/run_lapisgs_benchmark.py" +# "/root/autodl-tmp/envs/mip-splatting scripts/run_mipsplatting_benchmark.py" +# --- Apr 13 --- +# "/root/autodl-tmp/envs/c3dgs scripts/run_c3dgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/compgs scripts/run_compgs.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/sogs scripts/run_sogs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/compact3dgs scripts/run_compact3dgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/context_gs scripts/run_contextgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/spec_gaussian scripts/run_specgaussian_benchmark.py" # 重，跳过 +# "/root/autodl-tmp/envs/fds_gs scripts/run_fdsgs_benchmark.py" # 已OOM，跳过 +# "/root/autodl-tmp/envs/cdc_gs scripts/run_cdcgs_benchmark.py" +# "/root/autodl-tmp/envs/ngs scripts/run_ngs_benchmark.py" +# "/root/autodl-tmp/envs/aaa_gaussians scripts/run_aaags_benchmark.py" +# "/root/autodl-tmp/envs/sort_free_gs scripts/run_sortfreegs_benchmark.py" +# "/root/autodl-tmp/envs/analytic_splatting scripts/run_analyticsplatting_benchmark.py" +# "/root/autodl-tmp/envs/rade_gs scripts/run_radegs_benchmark.py" +# "/root/autodl-tmp/envs/splatfacto_w scripts/run_splatfactow_benchmark.py" # nerfstudio体系，结构不同，单独跑 +# "/root/autodl-tmp/envs/sundae scripts/run_sundae_benchmark.py" # 重，跳过 +# "/root/autodl-tmp/envs/pup_3dgs scripts/run_pup3dgs_benchmark.py" # 压缩后处理，跳过 +# "/root/autodl-tmp/envs/improving_adc scripts/run_improvingadc_benchmark.py" +# "/root/autodl-tmp/envs/topology_gs scripts/run_topologygs_benchmark.py" +# "/root/autodl-tmp/envs/GeoGaussian scripts/run_geogaussian_benchmark.py" # 重，跳过 +# "/root/autodl-tmp/envs/FLoD scripts/run_flod_benchmark.py" # 多LOD重，跳过 +# "/root/autodl-tmp/envs/VDGS scripts/run_vdgs_benchmark.py" +# # # --- Apr 14 --- +# # "/root/autodl-tmp/envs/gspull scripts/run_gspull_benchmark.py" +# "/root/autodl-tmp/envs/gaussianshader scripts/run_gsshader_benchmark.py" +# "/root/autodl-tmp/envs/gaussian_surfels scripts/run_surfel_benchmark.py" # 重，跳过 +"/root/autodl-tmp/envs/gaussianfocus scripts/run_gaussianfocus_benchmark.py" +"/root/autodl-tmp/envs/gigs scripts/run_gigs_benchmark.py" +"/root/autodl-tmp/envs/bags scripts/run_bags_benchmark.py" +"/root/autodl-tmp/envs/drk scripts/run_drkgs_benchmark.py" +# "/root/autodl-tmp/envs/fdsgs scripts/run_fds_benchmark.py" # 已OOM +"/root/autodl-tmp/envs/op43dgs scripts/run_op43dgs_benchmark.py" +"/root/autodl-tmp/envs/lp_3dgs scripts/run_lp3dgs_benchmark.py" +) + +echo "SplatAtlas: Starting New Methods Pipeline..." +echo "Errors will be logged to: $ERROR_LOG" + +for task in "${tasks[@]}"; do + read -r env_path script_path <<< "$task" + method_name=$(basename "$script_path") + + echo "----------------------------------------------------------------" + echo "Method: $method_name" + + conda activate "$env_path" + + if [ $? -eq 0 ]; then + echo "Checking requirements in $env_path..." + pip install lpips imageio scikit-image -q + + echo "Environment ready. Executing..." + + PYTHONPATH=/root/autodl-tmp/SplatAtlas "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" 2> >(tee -a "$ERROR_LOG" >&2) + + if [ $? -eq 0 ]; then + echo "Finished: $method_name" + else + echo "Error: $method_name failed. Saved to $ERROR_LOG" + echo "[$(date)] FAILED: $method_name" >> "$ERROR_LOG" + fi + else + echo "Failed to activate: $env_path" + echo "[$(date)] FAILED (Env Activation): $env_path for $method_name" >> "$ERROR_LOG" + fi + + conda deactivate +done + +echo "================================================================" +echo "Pipeline completed. Check $ERROR_LOG for crash reports." diff --git a/scripts/run_ngs_benchmark.py b/scripts/run_ngs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..f78163c627a7fb1fd89494689307b316765c1819 --- /dev/null +++ b/scripts/run_ngs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/ngs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/ngs:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"ckpts/ckpt_{iters}_rank0.pt") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_octree_gs_benchmark.py b/scripts/run_octree_gs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..ca471c4e258ba663381ebe0452190a2705e0b4e4 --- /dev/null +++ b/scripts/run_octree_gs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] [{name}] Elapsed: {elapsed:.1f}s") + return False + print(f" [COMPLETED] [{name}] Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] [{name}]: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/octree_gs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = "/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Octree-GS_official" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Physics", physics_cmd, curr_env) + + run_step("Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_op43dgs_benchmark.py b/scripts/run_op43dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4607f66be4be3880340500a4735b8cbc4db64457 --- /dev/null +++ b/scripts/run_op43dgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/op43dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/OP43DGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_opti3dgs_benchmark.py b/scripts/run_opti3dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..2fec3ebd74035826819e178d46717d45b2709ec1 --- /dev/null +++ b/scripts/run_opti3dgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] [{name}] Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] [{name}] Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] [{name}] {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/opti3dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Opti3DGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[{scene}] Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_pgsr_benchmark.py b/scripts/run_pgsr_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..f6f7aed67d8253aff75945d77e36801808cd39cf --- /dev/null +++ b/scripts/run_pgsr_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [ERROR] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/pgsr_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/pgsr_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_pgsrsgf_tnt.py b/scripts/run_pgsrsgf_tnt.py new file mode 100644 index 0000000000000000000000000000000000000000..62a6d2b683e24803c53586bd50f2df5b8b0d8348 --- /dev/null +++ b/scripts/run_pgsrsgf_tnt.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/pgsrsgf_tnt_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/pgsr_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"Done: {scene}\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_pipeline.py b/scripts/run_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..9715de439a76b998a53505401a8a650297890c7a --- /dev/null +++ b/scripts/run_pipeline.py @@ -0,0 +1,65 @@ +import os +import yaml +import subprocess +import time + +def run_command(cmd): + print(f"\n[EXEC] {cmd}") + result = subprocess.run(cmd, shell=True) + return result.returncode == 0 + +def main(): + config_path = "configs/mass_production.yaml" + with open(config_path, "r") as f: + config = yaml.safe_load(f) + + iters = config["global"]["iterations"] + + # 强制将当前目录加入环境变量，防止模块找不到 + env_setup = "export PYTHONPATH=$PYTHONPATH:$(pwd) && " + + for task in config["tasks"]: + method = task["method"] + base_path = task["base_path"] + + for scene in task["scenes"]: + print(f"\n{'='*60}") + print(f"🚀 开始流水线处理: [{method}] - [{scene}]") + print(f"{'='*60}") + + source_path = os.path.join(base_path, scene) + model_path = f"outputs/{method}_{scene}" + + # 1. 启动训练 + train_cmd = f"{env_setup} python scripts/main_train.py --method {method} --source_path {source_path} --model_path {model_path} --iterations {iters} --track_decoupling" + if not run_command(train_cmd): + print(f"❌ [跳过] {scene} 训练失败，可能遇到 OOM 或数据缺失。") + continue + + # 2. 全维度资产渲染 (RGB, Depth, Normal, Ablation) + render_cmd = f"{env_setup} python scripts/main_render.py --method {method} --source_path {source_path} --model_path {model_path} --iteration {iters}" + if not run_command(render_cmd): + print(f"❌ [跳过] {scene} 渲染失败。") + continue + + # 3. 执行测试集评测 (PSNR, SSIM, LPIPS) + render_dir = f"{model_path}/renders_test_{iters}" + gt_dir = f"{model_path}/gt_test_{iters}" + ply_path = f"{model_path}/point_cloud/iteration_{iters}/point_cloud.ply" + out_json = f"{model_path}/metrics_test_iter{iters}.json" + + eval_cmd = f"{env_setup} python ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {render_dir} --gt_dir {gt_dir} --ply_path {ply_path} --output_json {out_json} --colmap_dir {source_path}" + run_command(eval_cmd) + + # 4. 离线物理属性解剖 (Gini, Condition, Billboard...) + physics_cmd = f"{env_setup} python scripts/compute_offline_physics.py" + run_command(physics_cmd) + + # 5. 更新总榜单 CSV + table_cmd = f"{env_setup} python scripts/build_grand_table.py" + run_command(table_cmd) + + print(f"✅✅✅ {scene} 场景全流程处理完毕！\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_pixelgs_benchmark.py b/scripts/run_pixelgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..8aa496728eaafba83fa5439cc84ba73bc9b9f016 --- /dev/null +++ b/scripts/run_pixelgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" ❌ [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" ✨ [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" ❌ [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/pixelgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Pixel-GS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n🔥 [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" ✅ [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" ⚠️ [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" ✅ [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" ✅ [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"✅✅✅ {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_pixelgssgf.py b/scripts/run_pixelgssgf.py new file mode 100644 index 0000000000000000000000000000000000000000..89f8c09044411ba50262c77f10db8e5141616f77 --- /dev/null +++ b/scripts/run_pixelgssgf.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f"[Step] {name}") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f"[Failed] {name} Elapsed: {elapsed:.1f}s") + return False + print(f"[Completed] {name} Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f"[Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/pixelgssgf_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Pixel-GS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"Dataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"Processing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print("Checkpoint Ply file found.") + if not os.path.exists(render_flag): + print("Cleanup Previous render incomplete, re-rendering.") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print("Checkpoint Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print("Checkpoint Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done.") + +if __name__ == "__main__": + main() diff --git a/scripts/run_pup3dgs_benchmark.py b/scripts/run_pup3dgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..e27a9fad755bc7134fde0c6092b0bef9ec0c37cd --- /dev/null +++ b/scripts/run_pup3dgs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/pup3dgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/PUP3DGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_qgs_benchmark.py b/scripts/run_qgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..4c2941549f1348cc0e05bebcc3b52a491ad76edd --- /dev/null +++ b/scripts/run_qgs_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/qgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/quadratic_gs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene} r={res}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Table Build", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_r3dg_benchmark.py b/scripts/run_r3dg_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..2fd88679945f4589b8418a33ef32ddf022ca85be --- /dev/null +++ b/scripts/run_r3dg_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/r3dg_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Relightable3DGaussian:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_radegs_benchmark.py b/scripts/run_radegs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0fa2257e5a91090528c28e44892835d2c4daee08 --- /dev/null +++ b/scripts/run_radegs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/radegs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/radegs:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_rdogaussian_benchmark.py b/scripts/run_rdogaussian_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..8a14a47f3bbfb0f8ca12f4ec08b9c473acbf4912 --- /dev/null +++ b/scripts/run_rdogaussian_benchmark.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] [{name}] Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] [{name}] Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] [{name}] {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/rdogaussian_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/RDO-Gaussian_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_reactgs_benchmark.py b/scripts/run_reactgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..1d2fd10c52d9eaa5f42420a1759a65f1d086ee45 --- /dev/null +++ b/scripts/run_reactgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [ERROR] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/reactgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/ReAct-GS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"--- {scene} Done! ---\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_scaffoldgs_benchmark.py b/scripts/run_scaffoldgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..1d653cfa9364beaa70749ffb7808c3a4b7d7f608 --- /dev/null +++ b/scripts/run_scaffoldgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/scaffoldgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/scaffold_gs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_sgf_sensitivity.py b/scripts/run_sgf_sensitivity.py new file mode 100644 index 0000000000000000000000000000000000000000..68aea1775c9d9994683d71cd9fc86bcc383bf81f --- /dev/null +++ b/scripts/run_sgf_sensitivity.py @@ -0,0 +1,134 @@ +"""SGF hyperparameter sensitivity: 7 configs × 6 scenes.""" +import os +import subprocess +import time +import json +from pathlib import Path + +ATLAS = "/root/autodl-tmp/SplatAtlas" +METHOD = "vanillasgf_sens" +ITERS = 30000 + +SCENES = [ + {"name": "bicycle", "source": "/root/autodl-tmp/dataset/360/bicycle", "res": 4}, + {"name": "bonsai", "source": "/root/autodl-tmp/dataset/360/bonsai", "res": 2}, + {"name": "francis", "source": "/root/autodl-tmp/dataset/tnt/francis", "res": 2}, + {"name": "garden", "source": "/root/autodl-tmp/dataset/360/garden", "res": 4}, + {"name": "m60", "source": "/root/autodl-tmp/dataset/tnt/m60", "res": 2}, + {"name": "panther", "source": "/root/autodl-tmp/dataset/tnt/panther", "res": 2}, +] + +EXPERIMENTS = [ + ("aL1-0.3", {"SGF_ALPHA_L1": "0.3"}), + ("aL1-0.9", {"SGF_ALPHA_L1": "0.9"}), + ("lam-0.02", {"SGF_LAMBDA_STRUCT": "0.02"}), + ("lam-0.10", {"SGF_LAMBDA_STRUCT": "0.10"}), + ("gam-0.8", {"SGF_GAMMA": "0.8"}), + ("gam-2.0", {"SGF_GAMMA": "2.0"}), + ("beta-1.0", {"SGF_BETA": "1.0"}), +] + +def run_one(tag, env_overrides, scene_cfg): + scene = scene_cfg["name"] + source = scene_cfg["source"] + res = scene_cfg["res"] + model_out = f"{ATLAS}/outputs/sgf_sens_{tag}_{scene}" + ply_path = f"{model_out}/point_cloud/iteration_{ITERS}/point_cloud.ply" + metrics_json = f"{model_out}/metrics_test_iter{ITERS}.json" + + if os.path.exists(metrics_json): + return + + env = os.environ.copy() + env["PYTHONPATH"] = f"{ATLAS}:/root/autodl-tmp/3dgsAtlas_official:{env.get('PYTHONPATH','')}" + env.update(env_overrides) + + print(f"\n [{tag}] {scene} (r={res})") + + if not os.path.exists(ply_path): + cmd = (f"python {ATLAS}/scripts/main_train.py --method {METHOD} " + f"--source_path {source} --model_path {model_out} " + f"--iterations {ITERS} --resolution {res}") + r = subprocess.run(cmd, shell=True, env=env) + if r.returncode != 0: + print(f" [{tag}/{scene}] Train FAILED") + return + + renders_dir = f"{model_out}/renders_test_{ITERS}" + render_flag = f"{model_out}/render_complete_{ITERS}.flag" + if not os.path.exists(render_flag): + cmd = (f"python {ATLAS}/scripts/main_render.py --method {METHOD} " + f"--source_path {source} --model_path {model_out} " + f"--iteration {ITERS} --resolution {res}") + subprocess.run(cmd, shell=True, env=env) + + if not os.path.exists(metrics_json): + cmd = (f"python {ATLAS}/ufd_evalkit/run_eval.py --method {METHOD} " + f"--render_dir {renders_dir} --gt_dir {model_out}/gt_test_{ITERS} " + f"--ply_path {ply_path} --output_json {metrics_json} " + f"--colmap_dir {source} --depth_dir {model_out}/depths_test_{ITERS}") + subprocess.run(cmd, shell=True, env=env) + +def collect(): + print(f"\n{'='*80}") + print(f" SGF Sensitivity — 6-scene averages") + print(f"{'='*80}") + print(f"{'Tag':>12} {'bicycle':>8} {'bonsai':>8} {'francis':>8} {'garden':>8} {'m60':>8} {'panther':>8} {'AVG':>8}") + print("-"*76) + + for tag, _ in EXPERIMENTS: + vals = [] + cells = [] + for sc in SCENES: + f = Path(f"{ATLAS}/outputs/sgf_sens_{tag}_{sc['name']}/metrics_test_iter{ITERS}.json") + if f.exists(): + d = json.loads(f.read_text()) + p = d.get("photometric", d) + psnr = p.get("PSNR", p.get("psnr", None)) + if psnr is not None: + vals.append(psnr) + cells.append(f"{psnr:>8.2f}") + else: + cells.append(f"{'?':>8}") + else: + cells.append(f"{'—':>8}") + avg = f"{sum(vals)/len(vals):>8.2f}" if vals else f"{'?':>8}" + print(f"{tag:>12} {' '.join(cells)} {avg}") + + print("-"*76) + print(f"{'default†':>12}", end="") + dvals = [] + for sc in SCENES: + f = Path(f"{ATLAS}/outputs/vanillasgf_{sc['name']}/metrics_test_iter{ITERS}.json") + if f.exists(): + d = json.loads(f.read_text()) + p = d.get("photometric", d) + psnr = p.get("PSNR", p.get("psnr", None)) + if psnr is not None: + dvals.append(psnr) + print(f" {psnr:>8.2f}", end="") + else: + print(f" {'?':>8}", end="") + else: + print(f" {'—':>8}", end="") + davg = f"{sum(dvals)/len(dvals):>8.2f}" if dvals else "?" + print(f" {davg}") + +def main(): + total = len(EXPERIMENTS) * len(SCENES) + done = 0 + for tag, overrides in EXPERIMENTS: + print(f"\n{'='*60}") + print(f" Config: {tag} | {overrides}") + print(f"{'='*60}") + for sc in SCENES: + done += 1 + mj = Path(f"{ATLAS}/outputs/sgf_sens_{tag}_{sc['name']}/metrics_test_iter{ITERS}.json") + status = "SKIP" if mj.exists() else "RUN" + print(f" [{done}/{total}] {tag}/{sc['name']} -> {status}") + run_one(tag, overrides, sc) + + collect() + +if __name__ == "__main__": + main() diff --git a/scripts/run_sgf_sequence.sh b/scripts/run_sgf_sequence.sh new file mode 100644 index 0000000000000000000000000000000000000000..adeed6d24b4689e699a69154e5e267a1e3733b2e --- /dev/null +++ b/scripts/run_sgf_sequence.sh @@ -0,0 +1,51 @@ +#!/bin/bash + +CONDA_BASE=$(conda info --base) +source "$CONDA_BASE/etc/profile.d/conda.sh" + +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +declare -a tasks=( + + "/root/autodl-tmp/envs/absgs_env scripts/run_absgssgf.py" + + "/root/autodl-tmp/envs/mini_splatting scripts/run_minisplattingsgf.py" + + "/root/autodl-tmp/envs/mip-splatting scripts/run_mipsplatting_benchmark.py" + "/root/autodl-tmp/envs/mip-splatting scripts/run_mipsplattingsgf.py" + "/root/autodl-tmp/envs/mini_splatting scripts/run_minisplatting_benchmark.py" + "/root/autodl-tmp/envs/pixelgs scripts/run_pixelgs_benchmark.py" + "/root/autodl-tmp/envs/pixelgs scripts/run_pixelgssgf.py" +) + +echo "Starting Pipeline..." + +for task in "${tasks[@]}"; do + read -r env_path script_path <<< "$task" + + echo "----------------------------------------------------------------" + echo "Method: $(basename "$script_path")" + + conda activate "$env_path" + + if [ $? -eq 0 ]; then + echo "Checking requirements in $env_path..." + pip install lpips imageio scikit-image -q + + echo "Environment ready. Executing..." + "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" + + if [ $? -eq 0 ]; then + echo "Finished: $(basename "$script_path")" + else + echo "Error: $(basename "$script_path") failed." + fi + else + echo "Failed to activate: $env_path" + fi + + conda deactivate +done + +echo "================================================================" +echo "Optimization Pipeline completed." diff --git a/scripts/run_sgfgs_benchmark.py b/scripts/run_sgfgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..da0726b1ccab4ae83f5891ed2944a67bf1e1654f --- /dev/null +++ b/scripts/run_sgfgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" ❌ [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" ✨ [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" ❌ [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/sgfgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/sgfgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n🔥 [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" ✅ [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" ⚠️ [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" ✅ [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" ✅ [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"✅✅✅ {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_sogs_benchmark.py b/scripts/run_sogs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..2d344a85a61cc54730b8762837fea65b628b4c7c --- /dev/null +++ b/scripts/run_sogs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Success] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/sogs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/sogs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[{scene}] Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_sortfreegs_benchmark.py b/scripts/run_sortfreegs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..a065c420cbaa77545db8f2f0864b74ee503df8ea --- /dev/null +++ b/scripts/run_sortfreegs_benchmark.py @@ -0,0 +1,80 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/sortfreegs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/sortfreegs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" Ply file found.") + if not os.path.exists(render_flag): + print(f" Previous render incomplete") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics", physics_cmd, curr_env) + + run_step("Updating Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"Done") + +if __name__ == "__main__": + main() diff --git a/scripts/run_specgaussian_benchmark.py b/scripts/run_specgaussian_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..1007f3134369040937495b8d1d13e9f12a9daf02 --- /dev/null +++ b/scripts/run_specgaussian_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/specgaussian_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/specgaussian_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" Checkpoint Ply file found.") + if not os.path.exists(render_flag): + print(f" Cleanup Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" Checkpoint Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" Checkpoint Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_splatfactow_benchmark.py b/scripts/run_splatfactow_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..de24fdc765cd446951e41197dd80f4ddb1b54679 --- /dev/null +++ b/scripts/run_splatfactow_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/splatfactow_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/splatfactow_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_steepgs.py b/scripts/run_steepgs.py new file mode 100644 index 0000000000000000000000000000000000000000..f85f870ab07c9f3092fc661b5a0e68b0619439f0 --- /dev/null +++ b/scripts/run_steepgs.py @@ -0,0 +1,73 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/steepgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/SteepGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n==== Dataset Domain: {ds['name']} ====") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[Processing Scene] {scene} (resolution={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_sugar_benchmark.py b/scripts/run_sugar_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..3f8b99c7ed65cef6600251688bd703a468539119 --- /dev/null +++ b/scripts/run_sugar_benchmark.py @@ -0,0 +1,85 @@ +import os +import yaml +import subprocess +import time +import sys +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/sugar_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + base_atlas = "/root/autodl-tmp/SplatAtlas" + base_sugar = "/root/autodl-tmp/sugar_official" + + clean_env = os.environ.copy() + clean_env["PYTHONPATH"] = f"{base_atlas}:{base_sugar}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[SUGAR] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path) or os.path.exists(os.path.join(model_out, f"sugar_{iters}.pt")): + print(f" [Checkpoint] Model file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"{sys.executable} {base_atlas}/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, clean_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"{sys.executable} {base_atlas}/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, clean_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"{sys.executable} {base_atlas}/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, clean_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"{sys.executable} {base_atlas}/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, clean_env) + + run_step("Updating Grand Table", f"{sys.executable} {base_atlas}/scripts/build_grand_table.py", clean_env) + + print(f"Scene {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_sundae_benchmark.py b/scripts/run_sundae_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..0e22e9aae6f5136da32d901f37f2045d091d46a9 --- /dev/null +++ b/scripts/run_sundae_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/sundae_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/SUNDAE_offy:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_surfel.py b/scripts/run_surfel.py new file mode 100644 index 0000000000000000000000000000000000000000..03a6078d7534a5b2b881f6d3a87cbd3822e7cee5 --- /dev/null +++ b/scripts/run_surfel.py @@ -0,0 +1,6 @@ +import subprocess +import os +env = os.environ.copy() +env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gaussian_surfels:{env.get('PYTHONPATH', '')}" +cmd = "python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method gaussian_surfel --source_path /root/autodl-tmp/dataset/nerf_synthetic/lego --model_path /root/autodl-tmp/SplatAtlas/outputs/surfel_lego --iterations 15000 --track_decoupling" +subprocess.run(cmd, shell=True, env=env) diff --git a/scripts/run_surfel_benchmark.py b/scripts/run_surfel_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..be55f19ee9d1fc71e1ebc4ded2b85b012652aed4 --- /dev/null +++ b/scripts/run_surfel_benchmark.py @@ -0,0 +1,87 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/surfel_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + cap = cfg["global_settings"].get("cap_gaussians", None) + cap_flag = f" --cap_gaussians {cap}" if cap else "" + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/gaussian_surfels:{curr_env.get('PYTHONPATH', '')}" + curr_env["LD_LIBRARY_PATH"] = "/root/autodl-tmp/envs/gaussian_surfels/lib:" + curr_env.get("LD_LIBRARY_PATH", "") + curr_env["TORCH_CUDA_ARCH_LIST"] = "8.9" + curr_env["CUDA_HOME"] = "/usr/local/cuda" + curr_env["CUDA_VISIBLE_DEVICES"] = curr_env.get("CUDA_VISIBLE_DEVICES") or "0,1" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(" [Checkpoint] Render assets complete.") + else: + train_cmd = f"/root/autodl-tmp/envs/gaussian_surfels/bin/python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling{cap_flag}" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"/root/autodl-tmp/envs/gaussian_surfels/bin/python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[Done] {scene} completed.\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_surfel_eval_4new_scenes.sh b/scripts/run_surfel_eval_4new_scenes.sh new file mode 100644 index 0000000000000000000000000000000000000000..8ada97abdbb3cdcd4731ae9a6383d60e5316d891 --- /dev/null +++ b/scripts/run_surfel_eval_4new_scenes.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +set -u +cd /root/autodl-tmp/SplatAtlas +OUT_ROOT="outputs/surfel_hypothesis_all" +LOG_DIR="${OUT_ROOT}/logs" +mkdir -p "$LOG_DIR" + +# 只跑 4 个新场景 +SCENES=("church" "courthouse" "ignatius" "meetingroom") +METHODS=( + "vanilla_3dgs" "3dgs_dr" "3dgsmcmc" "aaags" "absgs" "absgssgf" + "analyticsplatting" "atomgs" "cdcgs" "coadaptgs" "conegs" "drkgs" + "erankgs" "ges" "ghap" "gslpm" "hogs" "lapisgs" "lightgaussian" + "minisplatting" "minisplattingsgf" "mipsplatting" "opti3dgs" "pixelgs" + "pixelgssgf" "rdogaussian" "reactgs" "steepgs" "tamgs" "gaussianpro" + "gof" "pgsr" "gsdf" "gspull" "2dgs" "trimgs" "gaussian_surfel" + "scaffoldgs" "octree_gs" "lod_gs" "sugar" +) + +# 捞取特殊迭代次数的逻辑（和之前一样） +iteration_arg() { + local m="$1" + case "$m" in + aaags) echo "--iteration 5000" ;; + tamgs) echo "--iteration 20000" ;; + gaussian_surfel) echo "--iteration 10000" ;; + hogs) echo "--iteration 30000" ;; + gsdf) echo "--iteration 20000" ;; + gspull) echo "--iteration 5000" ;; + sugar) echo "--iteration 5000" ;; + *) echo "" ;; + esac +} + +echo "============================================================" +echo "Evaluating 4 NEW T&T Scenes for Surfel Hypothesis..." +echo "============================================================" + +for m in "${METHODS[@]}"; do + for s in "${SCENES[@]}"; do + itarg="$(iteration_arg "$m")" + log="${LOG_DIR}/${m}_${s}.log" + echo "Running: method=$m scene=$s" + + # shellcheck disable=SC2086 + python scripts/eval_surfel_hypothesis.py \ + --method "$m" \ + --scene "$s" \ + $itarg \ + --mode subsample \ + --n-sample 200000 \ + --distance-multiplier 2.0 \ + --gt-normal-k 30 \ + --recon-pca-k 30 \ + --max-normal-points 50000 \ + --seed 0 \ + --verbose > "$log" 2>&1 + + status=$? + if [ $status -eq 0 ]; then + echo "[OK] $m $s" + else + echo "[FAILED] $m $s (Check log for reason)" + fi + done +done +echo "============================================================" +echo "Finished evaluating 4 new scenes!" +echo "============================================================" diff --git a/scripts/run_surfel_hypothesis_all_usable.sh b/scripts/run_surfel_hypothesis_all_usable.sh new file mode 100644 index 0000000000000000000000000000000000000000..fcaaf87184f9ca151f081dd00c483385ba0b3581 --- /dev/null +++ b/scripts/run_surfel_hypothesis_all_usable.sh @@ -0,0 +1,192 @@ +#!/usr/bin/env bash +set -u + +cd /root/autodl-tmp/SplatAtlas + +OUT_ROOT="outputs/surfel_hypothesis_all" +LOG_DIR="${OUT_ROOT}/logs" +mkdir -p "$LOG_DIR" + +SCENES=("barn" "caterpillar" "truck") + +# 41 usable methods under current audit policy: +# - standard 3DGS / geometry-aware: V3 + PCA when fields support it +# - 2dgs: method-specific 2D disk normal +# - trimgs: PCA-only / special primitive +# - gaussian_surfel included as usable ordinary primitive unless later evidence says otherwise +# - scaffold/octree/lod/sugar included, with group caveats in aggregation +METHODS=( + "vanilla_3dgs" + + "3dgs_dr" + "3dgsmcmc" + "aaags" + "absgs" + "absgssgf" + "analyticsplatting" + "atomgs" + "cdcgs" + "coadaptgs" + "conegs" + "drkgs" + "erankgs" + "ges" + "ghap" + "gslpm" + "hogs" + "lapisgs" + "lightgaussian" + "minisplatting" + "minisplattingsgf" + "mipsplatting" + "opti3dgs" + "pixelgs" + "pixelgssgf" + "rdogaussian" + "reactgs" + "steepgs" + "tamgs" + + "gaussianpro" + "gof" + "pgsr" + "gsdf" + "gspull" + + "2dgs" + "trimgs" + "gaussian_surfel" + "scaffoldgs" + "octree_gs" + "lod_gs" + "sugar" +) + +iteration_arg() { + local m="$1" + + case "$m" in + aaags) + echo "--iteration 5000" + ;; + tamgs) + echo "--iteration 20000" + ;; + gaussian_surfel) + # Audit found 5000/10000, no 15000. Use max available. + echo "--iteration 10000" + ;; + hogs) + # Expected 50000, audit found up to 30000. Use max available. + echo "--iteration 30000" + ;; + gsdf) + # Audit found up to 20000. + echo "--iteration 20000" + ;; + gspull) + # Audit found only 5000. + echo "--iteration 5000" + ;; + sugar) + # Audit found only 5000. + echo "--iteration 5000" + ;; + *) + echo "" + ;; + esac +} + +group_name() { + local m="$1" + + case "$m" in + gof|pgsr|gaussianpro|gsdf|gspull) + echo "geometry_aware_standard_or_usable" + ;; + 2dgs) + echo "surface_primitive_2dgs_disk_normal" + ;; + trimgs) + echo "surface_primitive_pca_only" + ;; + gaussian_surfel) + echo "gaussian_surfel_usable" + ;; + scaffoldgs|octree_gs|lod_gs) + echo "hierarchical_or_anchor_usable" + ;; + sugar) + echo "surface_mesh_coupled_usable" + ;; + aaags|tamgs|hogs) + echo "standard_3dgs_nondefault_checkpoint" + ;; + *) + echo "standard_3dgs" + ;; + esac +} + +echo "============================================================" +echo "Surfel hypothesis full run" +echo "methods: ${#METHODS[@]}" +echo "scenes: ${#SCENES[@]}" +echo "output root: $OUT_ROOT" +echo "============================================================" + +RUN_INDEX="${OUT_ROOT}/run_index.csv" +echo "method,scene,group,iteration_arg,status,log_path,json_path" > "$RUN_INDEX" + +for m in "${METHODS[@]}"; do + for s in "${SCENES[@]}"; do + g="$(group_name "$m")" + itarg="$(iteration_arg "$m")" + log="${LOG_DIR}/${m}_${s}.log" + json_path="outputs/surfel_hypothesis/${m}_${s}/surfel_hypothesis_eval.json" + + echo "" + echo "============================================================" + echo "Running: method=$m scene=$s group=$g iteration_arg='${itarg}'" + echo "Log: $log" + echo "============================================================" + + rm -f "$json_path" + + # shellcheck disable=SC2086 + python scripts/eval_surfel_hypothesis.py \ + --method "$m" \ + --scene "$s" \ + $itarg \ + --mode subsample \ + --n-sample 200000 \ + --distance-multiplier 2.0 \ + --gt-normal-k 30 \ + --recon-pca-k 30 \ + --max-normal-points 50000 \ + --seed 0 \ + --verbose > "$log" 2>&1 + + status=$? + + if [ $status -eq 0 ] && [ -f "$json_path" ]; then + echo "[OK] $m $s" + echo "${m},${s},${g},\"${itarg}\",ok,${log},${json_path}" >> "$RUN_INDEX" + + tail -80 "$log" | grep -E \ + '"method"|"scene"|"surface_near_ratio"|"n_normal_eval"|"has_v3_normal"|"has_2dgs_normal"|"has_explicit_normal"|"center_pca_normal_error_median_deg"|"v3_normal_error_median_deg"|"two_dgs_normal_error_median_deg"|"median_delta_v3_minus_center_pca_deg"|"median_delta_2dgs_minus_center_pca_deg"|"cov_log_anisotropy_sigma_auroc_surface_near"|"cov_planarity_auroc_surface_near"|"opacity_sigmoid_auroc_surface_near"|"spatial_pca_planarity_auroc_surface_near"' \ + || true + else + echo "[FAILED] $m $s status=$status" + echo "${m},${s},${g},\"${itarg}\",failed,${log},${json_path}" >> "$RUN_INDEX" + tail -80 "$log" || true + fi + done +done + +echo "" +echo "============================================================" +echo "Raw run finished. Run index:" +echo "$RUN_INDEX" +echo "============================================================" diff --git a/scripts/run_tamgs_benchmark.py b/scripts/run_tamgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..6e1d0814260d0b45def2736c0aeda408d3bd3acf --- /dev/null +++ b/scripts/run_tamgs_benchmark.py @@ -0,0 +1,109 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAIL] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [DONE] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [FAIL] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/tamgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/taming-3dgs:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + print(f" [SKIP] Source not found: {source}") + continue + + print(f"\n [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [OK] Checkpoint: Ply file found.") + if not os.path.exists(render_flag): + print(f" [WARN] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): + shutil.rmtree(d_path) + + render_cmd = ( + f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py" + f" --method {method} --source_path {source}" + f" --model_path {model_out} --iteration {iters} --resolution {res}" + ) + if not run_step("Offline Rendering", render_cmd, curr_env): + continue + else: + print(f" [OK] Checkpoint: Render assets complete.") + else: + train_cmd = ( + f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py" + f" --method {method} --source_path {source}" + f" --model_path {model_out} --iterations {iters}" + f" --resolution {res} --track_decoupling" + ) + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): + continue + + if not os.path.exists(metrics_json): + eval_cmd = ( + f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py" + f" --method {method} --scene {scene}" + f" --render_dir {renders_dir}" + f" --gt_dir {model_out}/gt_test_{iters}" + f" --ply_path {ply_path}" + f" --output_json {metrics_json}" + f" --colmap_dir {source}" + f" --depth_dir {depths_dir}" + ) + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [OK] Checkpoint: Metrics JSON already exists.") + + physics_cmd = ( + f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py" + f" --method {method} --source_path {source}" + f" --model_path {model_out} --iteration {iters}" + ) + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", + "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f" [DONE] {scene} Done!\n") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/scripts/run_texturedgs_benchmark.py b/scripts/run_texturedgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..bba5ce7f495cdffa2bc630b0d1c5e4c1d01287b9 --- /dev/null +++ b/scripts/run_texturedgs_benchmark.py @@ -0,0 +1,65 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [ERROR] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/texturedgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/textured_gaussians:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n==== Dataset: {ds['name']} ====") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {model_out}/renders_test_{iters} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json}" + run_step("Evaluation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Physics Diagnosis", physics_cmd, curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_texturegs_benchmark.py b/scripts/run_texturegs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..8dbee8ad9ca5e015b0c747d2d77e7e5cd3757b6f --- /dev/null +++ b/scripts/run_texturegs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s") + return False + print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name}: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/texturegs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/Texture-GS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_texturegs_full.py b/scripts/run_texturegs_full.py new file mode 100644 index 0000000000000000000000000000000000000000..f4a1a4fc893265d3b3576b3c3eb776f07cc0d024 --- /dev/null +++ b/scripts/run_texturegs_full.py @@ -0,0 +1,10 @@ +import os +import subprocess + +method = "texture_gs" +iters = 75000 +source = "/root/autodl-tmp/dataset/dtu/scan114" +model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_scan114" + +train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution 1 --track_decoupling" +subprocess.run(train_cmd, shell=True) diff --git a/scripts/run_topologygs_benchmark.py b/scripts/run_topologygs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..59702f9266e6156dcd8ed2e4843e0ec8c2ea6626 --- /dev/null +++ b/scripts/run_topologygs_benchmark.py @@ -0,0 +1,49 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [FAILED] [{name}] Elapsed: {elapsed:.1f}s") + return False + print(f" [OK] [{name}] Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [ERROR] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/topologygs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/TopologyGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + + if not run_step("Training TopologyGS", train_cmd, curr_env): + continue + +if __name__ == "__main__": + main() diff --git a/scripts/run_trimgs_benchmark.py b/scripts/run_trimgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..6d9685ffd52a84ce0b7bf5c37b4eb293899c597e --- /dev/null +++ b/scripts/run_trimgs_benchmark.py @@ -0,0 +1,81 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [ERROR] [{name}] Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [SUCCESS] [{name}] Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [EXCEPTION] [{name}] Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/trimgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/TrimGS_official/Trim2DGS:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\n[INFO] [{method.upper()}] Processing Scene: {scene} (r={res})") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + print(f" [INFO] [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [CLEANUP] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [INFO] [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [INFO] [Checkpoint] Metrics JSON already exists.") + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + + print(f"[SUCCESS] {scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/run_vanilla_benchmark.py b/scripts/run_vanilla_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..9ba4cad9a2c938db4279637aded4172e52181f06 --- /dev/null +++ b/scripts/run_vanilla_benchmark.py @@ -0,0 +1,71 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/vanilla_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/3dgs_official:{curr_env.get('PYTHONPATH', '')}" + + for ds in cfg["datasets"]: + print(f"\nDataset Domain: {ds['name']}") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + + if not os.path.exists(source): + continue + + print(f"\nProcessing Scene: {scene}") + + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + + if os.path.exists(ply_path): + if not os.path.exists(render_flag): + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training", train_cmd, curr_env): continue + + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + +if __name__ == "__main__": + main() diff --git a/scripts/run_vdgs_benchmark.py b/scripts/run_vdgs_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..895f77dd29c38bdb572b6cb1df2d6e96ff80b3f1 --- /dev/null +++ b/scripts/run_vdgs_benchmark.py @@ -0,0 +1,69 @@ +import os +import yaml +import subprocess +import time +import shutil + +def run_step(name, cmd, env): + print(f" [Step] {name}...") + start_time = time.time() + try: + result = subprocess.run(cmd, shell=True, env=env) + elapsed = time.time() - start_time + if result.returncode != 0: + print(f" [Error] {name} Failed. Elapsed: {elapsed:.1f}s") + return False + print(f" [Success] {name} Completed. Elapsed: {elapsed:.1f}s") + return True + except Exception as e: + print(f" [Exception] {name} Exception: {e}") + return False + +def main(): + config_path = "/root/autodl-tmp/SplatAtlas/configs/vdgs_benchmark.yaml" + with open(config_path, "r") as f: + cfg = yaml.safe_load(f) + method = cfg["method_name"] + iters = cfg["global_settings"]["iterations"] + curr_env = os.environ.copy() + curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/vdgs:{curr_env.get('PYTHONPATH', '')}" + for ds in cfg["datasets"]: + print(f"\n>>>> Dataset Domain: {ds['name']} <<<<") + for scene in ds["scenes"]: + source = os.path.join(ds["base_path"], scene) + model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}" + res = ds.get("resolution", 1) + if not os.path.exists(source): + continue + print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})") + ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply") + metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json") + render_flag = os.path.join(model_out, f"render_complete_{iters}.flag") + renders_dir = os.path.join(model_out, f"renders_test_{iters}") + depths_dir = os.path.join(model_out, f"depths_test_{iters}") + if os.path.exists(ply_path): + print(f" [Checkpoint] Ply file found.") + if not os.path.exists(render_flag): + print(f" [Cleanup] Previous render incomplete, re-rendering...") + for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]: + d_path = os.path.join(model_out, f"{d}_{iters}") + if os.path.exists(d_path): shutil.rmtree(d_path) + render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}" + if not run_step("Offline Rendering", render_cmd, curr_env): continue + else: + print(f" [Checkpoint] Render assets complete.") + else: + train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling" + if not run_step("Training & In-situ Rendering", train_cmd, curr_env): continue + if not os.path.exists(metrics_json): + eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}" + run_step("Metrics Calculation", eval_cmd, curr_env) + else: + print(f" [Checkpoint] Metrics JSON already exists.") + physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}" + run_step("Offline Physics Diagnosis", physics_cmd, curr_env) + run_step("Updating Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env) + print(f"{scene} Done!\n") + +if __name__ == "__main__": + main() diff --git a/scripts/train_missing_direct.sh b/scripts/train_missing_direct.sh new file mode 100644 index 0000000000000000000000000000000000000000..2069ed6a6233e9f79aa36f18f306c8c18ce074f4 --- /dev/null +++ b/scripts/train_missing_direct.sh @@ -0,0 +1,117 @@ + +#!/bin/bash + +set +e + +source /root/miniconda3/etc/profile.d/conda.sh + + + +ERROR_LOG="/root/autodl-tmp/SplatAtlas/patch_train_errors_direct.log" + +echo "================================================================" > "$ERROR_LOG" + +echo "Missing Scenes Patch Run Started at $(date)" >> "$ERROR_LOG" + +echo "================================================================" >> "$ERROR_LOG" + + + +declare -a tasks=( + + "/root/autodl-tmp/envs/aaa_gaussians scripts/run_aaags_benchmark.py" + + "/root/autodl-tmp/envs/absgs_env scripts/run_absgssgf_benchmark.py" + + "/root/autodl-tmp/envs/analytic_splatting scripts/run_analyticsplatting_benchmark.py" + + "/root/autodl-tmp/envs/atomgs scripts/run_atomgs_benchmark.py" + + "/root/autodl-tmp/envs/mini_splatting scripts/run_minisplattingsgf_benchmark.py" + + "/root/autodl-tmp/envs/pgsr_v2 scripts/run_pgsr_benchmark.py" + + "/root/autodl-tmp/envs/pixelgs scripts/run_pixelgssgf_benchmark.py" + + "/root/autodl-tmp/envs/scaffold_gs scripts/run_scaffoldgs_benchmark.py" + + "/root/autodl-tmp/envs/steepgs_env scripts/run_steepgs.py" + + "/root/autodl-tmp/env scripts/run_sugar_benchmark.py" + + "/root/autodl-tmp/envs/taming_3dgs scripts/run_tamgs_benchmark.py" + + "base scripts/run_vanilla_benchmark.py" + +) + + + +echo "SplatAtlas: Starting Missing Patches Pipeline..." + +echo "Errors will be logged to: $ERROR_LOG" + + + +for task in "${tasks[@]}"; do + + read -r env_path script_path <<< "$task" + + method_name=$(basename "$script_path") + + + + echo "----------------------------------------------------------------" + + echo "Method: $method_name" + + conda activate "$env_path" + + + + if [ $? -eq 0 ]; then + + echo "Checking requirements in $env_path..." + + pip install lpips imageio scikit-image -q + + echo "Environment ready. Executing..." + + + + # 核心修复点：仅在当前行提供干净的 PYTHONPATH 隔离运行，防止 arguments 模块串台 + + PYTHONPATH=/root/autodl-tmp/SplatAtlas "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" 2> >(tee -a "$ERROR_LOG" >&2) + + + + if [ $? -eq 0 ]; then + + echo "Finished: $method_name" + + else + + echo "Error: $method_name failed. Saved to $ERROR_LOG" + + echo "[$(date)] FAILED: $method_name" >> "$ERROR_LOG" + + fi + + else + + echo "Failed to activate: $env_path" + + echo "[$(date)] FAILED (Env Activation): $env_path for $method_name" >> "$ERROR_LOG" + + fi + + conda deactivate + +done + + + +echo "================================================================" + +echo "Pipeline completed. Check $ERROR_LOG for crash reports." + diff --git a/scripts/train_missing_patches.sh b/scripts/train_missing_patches.sh new file mode 100644 index 0000000000000000000000000000000000000000..475b1df2517b2936719a53aed007567f7351e4bf --- /dev/null +++ b/scripts/train_missing_patches.sh @@ -0,0 +1,36 @@ +#!/bin/bash +source /root/miniconda3/etc/profile.d/conda.sh +export PYTHONPATH=$PYTHONPATH:/root/autodl-tmp/SplatAtlas + +ERROR_LOG="/root/autodl-tmp/SplatAtlas/patch_train_errors.log" +echo "==== 启动缺失场景补训 ====" > "$ERROR_LOG" + +declare -a tasks=( + "aaags|/root/autodl-tmp/envs/aaa_gaussians|scripts/run_aaags_benchmark.py|configs/patch_aaags_benchmark.yaml" + "analyticsplatting|/root/autodl-tmp/envs/analytic_splatting|scripts/run_analyticsplatting_benchmark.py|configs/patch_analyticsplatting_benchmark.yaml" + "atomgs|/root/autodl-tmp/envs/atomgs|scripts/run_atomgs_benchmark.py|configs/patch_atomgs_benchmark.yaml" + "scaffoldgs|/root/autodl-tmp/envs/scaffold_gs|scripts/run_scaffoldgs_benchmark.py|configs/patch_scaffoldgs_benchmark.yaml" + "steepgs|/root/autodl-tmp/envs/steepgs_env|scripts/run_steepgs.py|configs/patch_steepgs_benchmark.yaml" + "tamgs|/root/autodl-tmp/envs/taming_3dgs|scripts/run_tamgs_benchmark.py|configs/patch_tamgs_benchmark.yaml" + "vanilla_3dgs|base|scripts/run_vanilla_benchmark.py|configs/patch_vanilla_benchmark.yaml" +) + +for task in "${tasks[@]}"; do + IFS='|' read -r method env_path script_path config_path <<< "$task" + echo "----------------------------------------------------------------" + echo ">> 正在补训: $method" + echo ">> 配置文件: $config_path" + + conda activate "$env_path" + if [ $? -eq 0 ]; then + python "$script_path" --config "$config_path" + if [ $? -ne 0 ]; then + echo "[FAILED] $method 训练出错" | tee -a "$ERROR_LOG" + fi + else + echo "[FAILED] 无法激活环境 $env_path" | tee -a "$ERROR_LOG" + fi + conda deactivate +done + +echo "==== 补训队列执行完毕 ====" diff --git a/scripts/train_selected_patches.sh b/scripts/train_selected_patches.sh new file mode 100644 index 0000000000000000000000000000000000000000..050248ef467278c047fd25ba633d2d10867cb06c --- /dev/null +++ b/scripts/train_selected_patches.sh @@ -0,0 +1,51 @@ +#!/bin/bash +set +e +source /root/miniconda3/etc/profile.d/conda.sh + +ERROR_LOG="/root/autodl-tmp/SplatAtlas/patch_train_errors_selected.log" +echo "================================================================" > "$ERROR_LOG" +echo "Selected Scenes Patch Run Started at $(date)" >> "$ERROR_LOG" +echo "================================================================" >> "$ERROR_LOG" + +declare -a tasks=( + "/root/autodl-tmp/envs/analytic_splatting scripts/run_analyticsplatting_benchmark.py configs/patch_analyticsplatting_benchmark.yaml" + "/root/autodl-tmp/envs/rdo-gaussian scripts/run_rdogaussian_benchmark.py configs/patch_rdogaussian_benchmark.yaml" + "/root/autodl-tmp/envs/steepgs_env scripts/run_steepgs.py configs/patch_steepgs_benchmark.yaml" + "/root/autodl-tmp/envs/gsdf scripts/run_gsdf_benchmark.py configs/patch_gsdf_benchmark.yaml" + "/root/autodl-tmp/env scripts/run_sugar_benchmark.py configs/patch_sugar_benchmark.yaml" +) + +echo "SplatAtlas: Starting Selected Patches Pipeline..." +echo "Errors will be logged to: $ERROR_LOG" + +for task in "${tasks[@]}"; do + read -r env_path script_path config_path <<< "$task" + method_name=$(basename "$script_path") + + echo "----------------------------------------------------------------" + echo ">> 正在补训: $method_name" + echo ">> 配置文件: $config_path" + + conda activate "$env_path" + + if [ $? -eq 0 ]; then + echo "Environment ready. Executing..." + + # 核心隔离运行逻辑 + PYTHONPATH=/root/autodl-tmp/SplatAtlas "$CONDA_PREFIX/bin/python" "/root/autodl-tmp/SplatAtlas/$script_path" --config "$config_path" 2> >(tee -a "$ERROR_LOG" >&2) + + if [ $? -eq 0 ]; then + echo "Finished: $method_name" + else + echo "Error: $method_name failed. Saved to $ERROR_LOG" + echo "[$(date)] FAILED: $method_name" >> "$ERROR_LOG" + fi + else + echo "Failed to activate: $env_path" + echo "[$(date)] FAILED (Env Activation): $env_path for $method_name" >> "$ERROR_LOG" + fi + conda deactivate +done + +echo "================================================================" +echo "Pipeline completed. Check $ERROR_LOG for crash reports." diff --git a/scripts/untitled.txt b/scripts/untitled.txt new file mode 100644 index 0000000000000000000000000000000000000000..816323184c34674d17db0b0a6aa086a335137b8c --- /dev/null +++ b/scripts/untitled.txt @@ -0,0 +1,9 @@ +# 设置倒计时秒数（2小时 = 7200秒） +seconds=7200; \ +while [ $seconds -gt 0 ]; do \ + echo -ne "倒计时: $(date -u -d @$seconds +%H:%M:%S)\r"; \ + sleep 1; \ + : $((seconds--)); \ +done; \ +echo -e "\n倒计时结束，开始运行脚本..."; \ +while true; do ./run_sgf_sequence.sh; done \ No newline at end of file diff --git a/scripts/verify_data.py b/scripts/verify_data.py new file mode 100644 index 0000000000000000000000000000000000000000..0d4b3588c598a4071a9605698495d83dd2178afe --- /dev/null +++ b/scripts/verify_data.py @@ -0,0 +1,31 @@ +import json +import os + +def check_metrics(): + target = "/root/autodl-tmp/SplatAtlas/outputs/2dgs_Lego/metrics_test_iter30000.json" + if not os.path.exists(target): + print("Metrics file not found. Ensure the pipeline ran successfully.") + return + + with open(target, 'r') as f: + data = json.load(f) + + essential_keys = [ + "PSNR", "SSIM", "LPIPS", + "gamma_median", "alpha_mean", "sh_energy_ratio", + "billboard_bias_ratio", "num_gaussians" + ] + + missing = [] + for k in essential_keys: + if k not in data or data[k] is None: + missing.append(k) + + if len(missing) == 0: + print("Validation Passed. All required physics, spatial, and dynamic decoupling metrics exist and are non-NaN.") + print(f"Extracted Sample -> PSNR: {data.get('PSNR', 'N/A')}, Billboard Bias: {data.get('billboard_bias_ratio', 'N/A')}") + else: + print(f"Validation Failed. Missing keys: {missing}") + +if __name__ == "__main__": + check_metrics() diff --git a/scripts/wrapper_generator_prompt.txt b/scripts/wrapper_generator_prompt.txt new file mode 100644 index 0000000000000000000000000000000000000000..b5416e5b4fd8478b179615c3dad8126a6a26cef4 --- /dev/null +++ b/scripts/wrapper_generator_prompt.txt @@ -0,0 +1,22 @@ +System Prompt for LLM to generate remaining 78 wrappers: + +I am building a 3DGS Benchmark. I have a BaseMethod class and a working Wrapper for Taming-3DGS (tamgs). +I need you to write a new wrapper for [INSERT_METHOD_NAME_HERE]. + +Rules: +1. Inherit from BaseMethod. +2. Register it with @register_method("[METHOD_NAME]"). +3. Implement `__init__`, `train_iteration`, `render`, `save`, `load_ply`. +4. The `train_iteration` MUST return a dictionary containing exactly: + - "loss" + - "num_gaussians" + - "pos_grad_norm" + - "gamma_median" + - "peak_vram_GB" + - "adc_trigger_rate" +5. If hyperparams.get("track_decoupling") is True, implement the Dual-Backward gradient decoupling exactly as done in the tamgs wrapper (calculate "grad_cos_sim" and "parasitic_ratio" at step % 100 == 0). +6. The `render` method MUST return a dict with "image" and, if available from the rasterizer, "depth". +7. Do not change the core algorithmic logic of [INSERT_METHOD_NAME_HERE], only adapt its inputs/outputs to match this exact interface. + +Reference Wrapper (TamGS): +[PASTE wrapper_tamgs.py HERE] diff --git a/tools/batch_render_coadaptgs_nonsynth.sh b/tools/batch_render_coadaptgs_nonsynth.sh new file mode 100644 index 0000000000000000000000000000000000000000..2f42ad7d38816b6cb14c689e0fe531721af9375b --- /dev/null +++ b/tools/batch_render_coadaptgs_nonsynth.sh @@ -0,0 +1,80 @@ +#!/bin/bash +set -u +ROOT="/root/autodl-tmp/SplatAtlas" +LOG="$ROOT/tools/batch_coadaptgs_render.log" +: > "$LOG" + +source /root/miniconda3/etc/profile.d/conda.sh +conda activate /root/autodl-tmp/envs/co_adaptation_3dgs +export PYTHONPATH=/root/autodl-tmp/CoAdaptGS:/root/autodl-tmp/SplatAtlas:$PYTHONPATH +cd "$ROOT" + +# scene source_path resolution +declare -a JOBS=( + # T&T (resolution 2) + "auditorium /root/autodl-tmp/dataset/tnt/auditorium 2" + "ballroom /root/autodl-tmp/dataset/tnt/ballroom 2" + "barn /root/autodl-tmp/dataset/tnt/barn 2" + "caterpillar /root/autodl-tmp/dataset/tnt/caterpillar 2" + "courtroom /root/autodl-tmp/dataset/tnt/courtroom 2" + "lighthouse /root/autodl-tmp/dataset/tnt/lighthouse 2" + "museum /root/autodl-tmp/dataset/tnt/museum 2" + "palace /root/autodl-tmp/dataset/tnt/palace 2" + "playground /root/autodl-tmp/dataset/tnt/playground 2" + "temple /root/autodl-tmp/dataset/tnt/temple 2" + "train /root/autodl-tmp/dataset/tnt/train 2" + "truck /root/autodl-tmp/dataset/tnt/truck 2" + # 360 outdoor (resolution 4) + "bicycle /root/autodl-tmp/dataset/360/bicycle 4" + "flowers /root/autodl-tmp/dataset/360/flowers 4" + "garden /root/autodl-tmp/dataset/360/garden 4" + "stump /root/autodl-tmp/dataset/360/stump 4" + "treehill /root/autodl-tmp/dataset/360/treehill 4" + # 360 indoor (resolution 2) + "bonsai /root/autodl-tmp/dataset/360/bonsai 2" + "counter /root/autodl-tmp/dataset/360/counter 2" + "kitchen /root/autodl-tmp/dataset/360/kitchen 2" + "room /root/autodl-tmp/dataset/360/room 2" + # Deep Blending (resolution 1) + "DrJohnson /root/autodl-tmp/dataset/deepblending_clean/DrJohnson 1" + "Playroom /root/autodl-tmp/dataset/deepblending_clean/Playroom 1" +) + +OK=0; FAIL=0; SKIP=0 +for job in "${JOBS[@]}"; do + read -r scene src res <<<"$job" + model="$ROOT/outputs/coadaptgs_$scene" + ply="$model/point_cloud/iteration_30000/point_cloud.ply" + + echo "" | tee -a "$LOG" + echo "================ $scene (res=$res) ================" | tee -a "$LOG" + + if [ ! -f "$ply" ]; then + echo "[SKIP] PLY not found: $ply" | tee -a "$LOG" + SKIP=$((SKIP+1)); continue + fi + if [ ! -d "$src" ]; then + echo "[SKIP] source not found: $src" | tee -a "$LOG" + SKIP=$((SKIP+1)); continue + fi + + python scripts/main_render.py \ + --method coadaptgs \ + --source_path "$src" \ + --model_path "$model" \ + --resolution "$res" >>"$LOG" 2>&1 + + rc=$? + rd_count=$(ls "$model/renders_test_30000"/*.png 2>/dev/null | wc -l) + gd_count=$(ls "$model/gt_test_30000"/*.png 2>/dev/null | wc -l) + if [ $rc -eq 0 ] && [ $rd_count -gt 0 ] && [ $gd_count -gt 0 ]; then + echo "[OK] $scene renders=$rd_count gt=$gd_count" | tee -a "$LOG" + OK=$((OK+1)) + else + echo "[FAIL] $scene rc=$rc renders=$rd_count gt=$gd_count" | tee -a "$LOG" + FAIL=$((FAIL+1)) + fi +done + +echo "" | tee -a "$LOG" +echo "=== SUMMARY: OK=$OK FAIL=$FAIL SKIP=$SKIP ===" | tee -a "$LOG" diff --git a/tools/batch_render_missing.py b/tools/batch_render_missing.py new file mode 100644 index 0000000000000000000000000000000000000000..e2f9bb751851b083539998e5075672fc92560b73 --- /dev/null +++ b/tools/batch_render_missing.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +import argparse +import os +import subprocess +from pathlib import Path + +ROOT = Path("/root/autodl-tmp/SplatAtlas") + +# 31 个需要补 render 的 cell +MISSING_CELLS = [ + ("coadaptgs", s) for s in [ + "Auditorium","Ballroom","Barn","Caterpillar","Courtroom","Lighthouse", + "Museum","Palace","Playground","Temple","Train","Truck", + "Bicycle","Bonsai","Counter","Flowers","Garden","Kitchen", + "Room","Stump","Treehill","DrJohnson","Playroom" + ] +] + [ + ("lod_gs", s) for s in ["Chair","Drums","Ficus","Hotdog","Lego","Materials","Mic","Ship"] +] + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--dry-run", action="store_true", help="只打印命令，不真正执行") + args = parser.parse_args() + + for method, scene in MISSING_CELLS: + run_dir = ROOT / "outputs" / f"{method}_{scene}" + if not run_dir.exists(): + print(f"SKIP {method}_{scene}: run dir 不存在") + continue + + # 构造和原脚本一样的调用 + cmd = [ + "python", "-m", "tools.render_offline", # ← 这里改成你实际的 render 入口 + "--method", method, + "--source_path", str(run_dir / "input"), # 大部分方法 source_path 在 input/ 里 + "--model_path", str(run_dir), + "--iteration", "30000" + ] + + print(f"\n=== {method} × {scene} ===") + print(" ".join(cmd)) + + if not args.dry_run: + try: + subprocess.run(cmd, check=True) + except subprocess.CalledProcessError as e: + print(f"❌ 失败: {e}") + +if __name__ == "__main__": + main() diff --git a/tools/build_and_compose_fig1.py b/tools/build_and_compose_fig1.py new file mode 100644 index 0000000000000000000000000000000000000000..e6cdd8557fffe39e78edaaecfa57321bb4c3780a --- /dev/null +++ b/tools/build_and_compose_fig1.py @@ -0,0 +1,229 @@ +import matplotlib.pyplot as plt +from PIL import Image, ImageDraw, ImageFont +import os +import glob +import json +import numpy as np +from pdf2image import convert_from_path + +# ========================================================================= +# 配置项 +# ========================================================================= +SCENE = "bonsai" +VIEW = "00000.png" +OUTPUTS_BASE = "/root/autodl-tmp/SplatAtlas/outputs" +METHODS = ["vanilla_3dgs", "pgsr", "ges", "opti3dgs", "steepgs"] + +LEFT_PDF_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf" +OUTPUT_PREFIX = "/root/autodl-tmp/SplatAtlas/tex/figures/fig1_saturation_puzzle" + +# ========================================================================= +# 辅助函数 +# ========================================================================= +def get_valid_dir(method, scene): + for suffix in ["", "_bak"]: + d = f"{OUTPUTS_BASE}/{method}_{scene}{suffix}" + if os.path.exists(d) and len(glob.glob(f"{d}/renders_test*")) > 0: + return d + return None + +def get_metrics(cell_dir): + try: + # 尝试读取 metrics json 以获取底栏数据 + metric_files = sorted(glob.glob(f"{cell_dir}/metrics_test_iter*.json")) + if not metric_files: return None + with open(metric_files[-1], 'r') as f: + data = json.load(f) + # 根据你之前的截图，提取 alpha_mean 和 gamma_median + # 如果你已经按我建议改成了倒数，请自行在此处调整计算 + geo = data.get("geometric", {}) + return { + "opacity": geo.get("alpha_mean", 0.0), + "aniso": geo.get("gamma_median", 0.0), # 如果是倒数可能是 condition_number 等 + "psnr": data.get("photometric", {}).get("PSNR", 0.0) + } + except Exception: + return None + +def add_label(img, title, psnr, opacity=None, aniso=None, is_gt=False): + """在图片上绘制带有半透明黑底的白字标签 (模拟截图效果)""" + draw = ImageDraw.Draw(img, 'RGBA') + w, h = img.size + + # 标题 (顶部居中) + title_text = title.upper() + font_size = max(16, int(h * 0.03)) + # 简单回退：如果找不到系统字体，使用默认字体。 + # 为了更好效果，这里尽量用比较大的字号，由于默认字体不可缩放，我们可能需要绘制矩形框来显眼 + try: + font = ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) + except: + font = ImageFont.load_default() + + # 标题底框 + padding = 5 + title_bbox = draw.textbbox((0, 0), title_text, font=font) + tw, th = title_bbox[2] - title_bbox[0], title_bbox[3] - title_bbox[1] + title_x = (w - tw) // 2 + title_y = 10 + draw.rectangle([title_x - padding, title_y - padding, + title_x + tw + padding, title_y + th + padding], + fill=(0, 0, 0, 180)) + # GT 用白色，Method 用绿色 + title_color = (255, 255, 255) if is_gt else (100, 255, 100) + draw.text((title_x, title_y), title_text, font=font, fill=title_color) + + # 底部 PSNR 标签 (右下) + if psnr: + psnr_text = f"PSNR: {psnr:.2f} dB" if not is_gt else "Reference" + p_bbox = draw.textbbox((0, 0), psnr_text, font=font) + pw, ph = p_bbox[2] - p_bbox[0], p_bbox[3] - p_bbox[1] + px = w - pw - 10 + py = h - ph - 40 # 留出底栏空间 + draw.rectangle([px - padding, py - padding, + px + pw + padding, py + ph + padding], + fill=(0, 0, 0, 180)) + draw.text((px, py), psnr_text, font=font, fill=(255,255,255)) + + # 底部黑条 (Opacity & Aniso) + if opacity is not None and aniso is not None: + bar_h = 30 + draw.rectangle([0, h - bar_h, w, h], fill=(30, 30, 30, 255)) + stat_text = f"opacity μ: {opacity:.3f} | anisotropy μ: {aniso:.3f}" if not is_gt else "opacity μ: 1.000 | anisotropy μ: 1.000" + s_bbox = draw.textbbox((0, 0), stat_text, font=font) + sw, sh = s_bbox[2] - s_bbox[0], s_bbox[3] - s_bbox[1] + draw.text(((w-sw)//2, h - bar_h + (bar_h-sh)//2), stat_text, font=font, fill=(200,200,200)) + + return img + + +# ========================================================================= +# 步骤 1: 构建右侧 6 宫格图片 +# ========================================================================= +print("正在构建右侧 6 宫格面板...") + +images = [] +labels = ["GT"] + METHODS + +# 提取 GT +gt_dir = get_valid_dir("vanilla_3dgs", SCENE) +if gt_dir: + gt_path = os.path.join(gt_dir, "gt_test_30000", VIEW) + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, "gt_test", VIEW) + + if os.path.exists(gt_path): + img_gt = Image.open(gt_path).convert("RGBA") + img_gt = add_label(img_gt, "GT", None, opacity=1.0, aniso=1.0, is_gt=True) + images.append(img_gt) + else: + print(f"❌ 找不到 GT 图像: {gt_path}") + exit(1) +else: + print("❌ 找不到 vanilla_3dgs 目录以提取 GT") + exit(1) + +# 提取 Methods +for method in METHODS: + cell_dir = get_valid_dir(method, SCENE) + if not cell_dir: + print(f"⚠️ 找不到方法目录: {method}") + # 创建占位黑图 + img_placeholder = Image.new("RGBA", images[0].size, (50,50,50,255)) + images.append(img_placeholder) + continue + + render_dirs = sorted(glob.glob(f"{cell_dir}/renders_test*")) + if not render_dirs: + print(f"⚠️ 找不到渲染目录: {method}") + images.append(Image.new("RGBA", images[0].size, (50,50,50,255))) + continue + + img_path = os.path.join(render_dirs[-1], VIEW) + if not os.path.exists(img_path): + print(f"⚠️ 找不到图像文件: {img_path}") + images.append(Image.new("RGBA", images[0].size, (50,50,50,255))) + continue + + img_m = Image.open(img_path).convert("RGBA") + + # 提取底栏数据 + metrics = get_metrics(cell_dir) + if metrics: + img_m = add_label(img_m, method, metrics['psnr'], metrics['opacity'], metrics['aniso']) + else: + img_m = add_label(img_m, method, 0.0) + + images.append(img_m) + +# 拼接 6 宫格 +iw, ih = images[0].size +padding = 10 +grid_w = iw * 3 + padding * 2 +grid_h = ih * 2 + padding + +img_right = Image.new("RGBA", (grid_w, grid_h), (30,30,30,255)) +positions = [ + (0, 0), (iw + padding, 0), ((iw + padding) * 2, 0), + (0, ih + padding), (iw + padding, ih + padding), ((iw + padding) * 2, ih + padding) +] + +for img, pos in zip(images, positions): + img_right.paste(img, pos) + + +# ========================================================================= +# 步骤 2: 与左侧 PDF 拼接 +# ========================================================================= +print("正在从 PDF 提取左侧面板 (DPI=300)...") +if not os.path.exists(LEFT_PDF_PATH): + print(f"❌ 找不到左图 (PDF): {LEFT_PDF_PATH}") + exit(1) + +try: + pages = convert_from_path(LEFT_PDF_PATH, dpi=300) + img_left = pages[0] +except Exception as e: + print(f"❌ PDF 转换失败: {e}\n请确保已安装 poppler-utils") + exit(1) + +print("正在拼接最终图像...") +target_height = 1500 +aspect_l = img_left.width / img_left.height +aspect_r = img_right.width / img_right.height + +new_w_l = int(target_height * aspect_l) +new_w_r = int(target_height * aspect_r) + +img_left = img_left.resize((new_w_l, target_height), Image.Resampling.LANCZOS) +img_right = img_right.resize((new_w_r, target_height), Image.Resampling.LANCZOS) + +fig, (ax1, ax2) = plt.subplots( + 1, 2, + figsize=(18, 6), + gridspec_kw={'width_ratios': [aspect_l, aspect_r]} +) + +ax1.imshow(img_left) +ax1.axis('off') +ax1.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", + fontsize=18, pad=12, loc='left', fontweight='bold') + +ax2.imshow(img_right) +ax2.axis('off') +ax2.set_title(f"(b) Saturated Cluster ({SCENE.upper()}): Render vs. Representation", + fontsize=18, pad=12, loc='left', fontweight='bold') + +plt.subplots_adjust(wspace=0.03) + +os.makedirs(os.path.dirname(OUTPUT_PREFIX), exist_ok=True) +pdf_path = f"{OUTPUT_PREFIX}.pdf" +png_path = f"{OUTPUT_PREFIX}.png" + +print("正在保存最终拼接文件...") +plt.savefig(pdf_path, dpi=300, bbox_inches='tight', format='pdf', transparent=True) +plt.savefig(png_path, dpi=300, bbox_inches='tight', format='png', transparent=False) + +print(f"=== Figure 1 (Saturation Puzzle) 拼图合成完毕 ===") +print(f"✅ Saved PDF: {pdf_path} (供 LaTeX 插入)") +print(f"✅ Saved PNG: {png_path} (供快速预览)") diff --git a/tools/build_and_compose_fig1_v3.py b/tools/build_and_compose_fig1_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..547badd9cb39568b8167079cdd1ff53741b33e5d --- /dev/null +++ b/tools/build_and_compose_fig1_v3.py @@ -0,0 +1,405 @@ +import os, glob +import numpy as np +import matplotlib.pyplot as plt +from PIL import Image, ImageDraw, ImageFont +from plyfile import PlyData + +try: + from pdf2image import convert_from_path + HAS_PDF2IMAGE = True +except ImportError: + HAS_PDF2IMAGE = False + +# ========================================================================= +# Config +# ========================================================================= +SCENE = 'truck' +REF_VIEW = '00014.png' # vanilla's view as reference for hash matching +OUTPUTS_BASE = '/root/autodl-tmp/SplatAtlas/outputs' +ASSETS_DIR = '/root/autodl-tmp/SplatAtlas/outputs/phase5b/gallery_assets_31' +LEFT_PDF = '/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf' +OUTPUT_PREFIX = '/root/autodl-tmp/SplatAtlas/tex/figures/fig1_saturation_puzzle' + +METHODS = ['minisplatting', '3dgsmcmc', 'gof', 'pixelgs', 'vanilla_3dgs'] + +# ========================================================================= +# Helpers +# ========================================================================= +def get_valid_dir(method, scene): + for suffix in ["", "_bak"]: + d = f"{OUTPUTS_BASE}/{method}_{scene}{suffix}" + if os.path.exists(d) and len(glob.glob(f"{d}/gt_test*")) > 0: + return d + return None + + +def find_matching_view(cell_dir, ref_arr): + """ + Hash-match: find the file in this method's gt_test/ that depicts + the same viewpoint as ref_arr. Returns (filename, gt_arr) or (None, None). + """ + gt_dirs = sorted(glob.glob(f"{cell_dir}/gt_test*")) + if not gt_dirs: + return None, None + + best_name, best_mse, best_arr = None, float('inf'), None + + for f in sorted(glob.glob(f"{gt_dirs[-1]}/*.png")): + try: + cand = np.array(Image.open(f).convert('RGB')).astype(np.float32) + + if cand.shape != ref_arr.shape: + continue + + diff = np.abs(ref_arr - cand) + diff[diff > 240] = 0 + diff[np.abs(diff - 128) < 15] = 0 + + mse = float(np.mean(diff ** 2)) + + if mse < best_mse: + best_mse, best_name, best_arr = mse, os.path.basename(f), cand + + except Exception: + continue + + return (best_name, best_arr) if best_name and best_mse < 50.0 else (None, None) + + +def get_render_and_psnr(cell_dir, view_name, gt_arr): + if view_name is None or gt_arr is None: + return None, None + + rdirs = sorted(glob.glob(f"{cell_dir}/renders_test*")) + if not rdirs: + return None, None + + rp = f"{rdirs[-1]}/{view_name}" + if not os.path.exists(rp): + return None, None + + try: + r_pil = Image.open(rp).convert('RGB') + r_arr = np.array(r_pil).astype(np.float32) + + if r_arr.shape != gt_arr.shape: + return None, None + + mse = float(np.mean((r_arr - gt_arr) ** 2)) + psnr = 100.0 if mse == 0 else 20 * np.log10(255.0 / np.sqrt(mse)) + + return r_pil, psnr + + except Exception: + return None, None + + +def get_ply_metrics(cell_dir): + ply_paths = sorted(glob.glob(f"{cell_dir}/point_cloud/iteration_*/point_cloud.ply")) + if not ply_paths: + return None + + try: + v = PlyData.read(ply_paths[-1])['vertex'] + + opacities = 1.0 / (1.0 + np.exp(-v['opacity'])) + + scales_log = np.vstack([ + v['scale_0'], + v['scale_1'], + v['scale_2'] + ]) + + scales_phys = np.exp(scales_log) + + log_aniso = np.max(scales_log, axis=0) - np.min(scales_log, axis=0) + + return { + 'alpha_med': float(np.median(opacities)), + 'scale_med': float(np.median(scales_phys)), + 'aniso_med': float(np.exp(np.median(log_aniso))), + 'n_gauss': int(scales_log.shape[1]), + } + + except Exception: + return None + + +def fmt_aniso(a): + return f"{a:.1f}" if a < 100 else f"{a:,.0f}" + + +def safe_load_font(font_name, size): + try: + return ImageFont.truetype(font_name, size) + except Exception: + return ImageFont.load_default() + + +# ========================================================================= +# Step 1: build right panel 3x2 grid via PIL +# ========================================================================= +print("[1/3] Building right panel (6 cells)...") + +ref_path = f"{ASSETS_DIR}/gt_{SCENE}_{REF_VIEW}" +if not os.path.exists(ref_path): + raise SystemExit(f"Missing reference GT: {ref_path}") + +ref_pil = Image.open(ref_path).convert('RGB') +ref_arr = np.array(ref_pil).astype(np.float32) + +W, H = ref_pil.size + +panels = [ + (ref_pil, 'GROUND TRUTH', None, None, True) +] + +for m in METHODS: + cell_dir = get_valid_dir(m, SCENE) + + if cell_dir is None: + panels.append((Image.new('RGB', (W, H), (50, 50, 50)), m, None, None, False)) + print(f" [{m}] MISSING_DIR") + continue + + view_name, gt_arr = find_matching_view(cell_dir, ref_arr) + img, psnr = get_render_and_psnr(cell_dir, view_name, gt_arr) + + if img is None or img.size != (W, H): + fallback = f"{ASSETS_DIR}/{m}_{SCENE}_{REF_VIEW}" + + if os.path.exists(fallback): + img = Image.open(fallback).convert('RGB') + + if img.size == (W, H): + img_arr = np.array(img).astype(np.float32) + mse = float(np.mean((img_arr - ref_arr) ** 2)) + psnr = 100.0 if mse == 0 else 20 * np.log10(255.0 / np.sqrt(mse)) + else: + img = Image.new('RGB', (W, H), (50, 50, 50)) + psnr = None + else: + img = Image.new('RGB', (W, H), (50, 50, 50)) + psnr = None + + metrics = get_ply_metrics(cell_dir) + panels.append((img, m, psnr, metrics, False)) + + if psnr is not None: + print(f" [{m}] view={view_name} PSNR={psnr:.2f}") + else: + print(f" [{m}] view={view_name}") + + +# ========================================================================= +# Layout +# ========================================================================= +gap = 12 +bottom_bar = int(H * 0.11) + +font_main = safe_load_font("DejaVuSans-Bold.ttf", int(H * 0.045)) +font_metric = safe_load_font("DejaVuSans-Bold.ttf", int(H * 0.030)) + +canvas_w = W * 3 + gap * 2 +canvas_h = (H + bottom_bar) * 2 + gap + +right_canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30)) +draw = ImageDraw.Draw(right_canvas) + +positions = [ + (0, 0), (1, 0), (2, 0), + (0, 1), (1, 1), (2, 1) +] + +for (img, name, psnr, metrics, is_gt), (col, row) in zip(panels, positions): + x = col * (W + gap) + y = row * (H + bottom_bar + gap) + + right_canvas.paste(img, (x, y)) + + pad = 8 + + # top-left method label + label = name.upper() + bb = draw.textbbox((0, 0), label, font=font_main) + lw, lh = bb[2] - bb[0], bb[3] - bb[1] + + draw.rectangle( + [x, y, x + lw + pad * 2, y + lh + pad * 2], + fill=(0, 0, 0) + ) + + color = (255, 255, 255) if is_gt else (76, 175, 80) + + draw.text( + (x + pad, y + pad), + label, + fill=color, + font=font_main + ) + + # bottom-right PSNR, overlaid on the image + if is_gt: + psnr_text = "Reference" + elif psnr is not None: + psnr_text = f"PSNR: {psnr:.2f} dB" + else: + psnr_text = "PSNR: ---" + + bb = draw.textbbox((0, 0), psnr_text, font=font_main) + pw, ph = bb[2] - bb[0], bb[3] - bb[1] + + draw.rectangle( + [x + W - pw - pad * 2, y + H - ph - pad * 2, x + W, y + H], + fill=(0, 0, 0) + ) + + draw.text( + (x + W - pw - pad, y + H - ph - pad), + psnr_text, + fill=(255, 255, 255), + font=font_main + ) + + # bottom bar metrics + if metrics is not None: + line1 = ( + f"Opacity μ: {metrics['alpha_med']:.3f} | " + f"Scale μ: {metrics['scale_med']:.4f} | " + f"N: {metrics['n_gauss'] / 1e6:.2f}M" + ) + + line2 = f"Anisotropy μ: {fmt_aniso(metrics['aniso_med'])}" + + b1 = draw.textbbox((0, 0), line1, font=font_metric) + b2 = draw.textbbox((0, 0), line2, font=font_metric) + + l1h = b1[3] - b1[1] + l2h = b2[3] - b2[1] + + total = l1h + l2h + 6 + + t1y = y + H + (bottom_bar - total) // 2 + t2y = t1y + l1h + 6 + + draw.text( + (x + (W - (b1[2] - b1[0])) // 2, t1y), + line1, + fill=(200, 200, 200), + font=font_metric + ) + + draw.text( + (x + (W - (b2[2] - b2[0])) // 2, t2y), + line2, + fill=(255, 152, 0), + font=font_metric + ) + + elif is_gt: + txt = "(reference)" + b = draw.textbbox((0, 0), txt, font=font_metric) + + draw.text( + ( + x + (W - (b[2] - b[0])) // 2, + y + H + (bottom_bar - (b[3] - b[1])) // 2 + ), + txt, + fill=(150, 150, 150), + font=font_metric + ) + + +# ========================================================================= +# Step 2: load left panel from PDF +# ========================================================================= +print("[2/3] Loading left panel from PDF...") + +if HAS_PDF2IMAGE and os.path.exists(LEFT_PDF): + try: + left_img = convert_from_path(LEFT_PDF, dpi=300)[0].convert('RGB') + except Exception as e: + print(f" WARNING: failed to load PDF with pdf2image: {e}") + print(" Placeholder used instead.") + left_img = Image.new('RGB', (1200, 900), (240, 240, 240)) +else: + print(f" WARNING: pdf2image missing or {LEFT_PDF} absent — placeholder used") + left_img = Image.new('RGB', (1200, 900), (240, 240, 240)) + + +# ========================================================================= +# Step 3: compose with matplotlib +# ========================================================================= +print("[3/3] Composing final figure...") + +aspect_l = left_img.width / left_img.height +aspect_r = right_canvas.width / right_canvas.height + +fig_w = 18.0 +fig_h = fig_w * 1.0 / (aspect_l + aspect_r + 0.05) + +fig = plt.figure(figsize=(fig_w, fig_h + 0.5)) + +gs = fig.add_gridspec( + 1, + 2, + width_ratios=[aspect_l, aspect_r], + wspace=0.05, + top=0.90, + bottom=0.02, + left=0.01, + right=0.99 +) + +ax1 = fig.add_subplot(gs[0]) +ax1.imshow(np.array(left_img)) +ax1.axis('off') + +ax2 = fig.add_subplot(gs[1]) +ax2.imshow(np.array(right_canvas)) +ax2.axis('off') + +ax1.text( + 0.0, + 1.015, + "(a) Per-seed PSNR rank instability · Bonsai & Lego", + transform=ax1.transAxes, + ha='left', + va='bottom', + fontsize=13, + fontweight='bold' +) + +ax2.text( + 0.0, + 1.015, + f"(b) Saturated cluster ({SCENE.title()}) · render vs. representation", + transform=ax2.transAxes, + ha='left', + va='bottom', + fontsize=13, + fontweight='bold' +) + +os.makedirs(os.path.dirname(OUTPUT_PREFIX), exist_ok=True) + +plt.savefig( + f"{OUTPUT_PREFIX}.pdf", + dpi=300, + bbox_inches='tight', + format='pdf' +) + +plt.savefig( + f"{OUTPUT_PREFIX}.png", + dpi=300, + bbox_inches='tight', + format='png', + facecolor='white' +) + +plt.close() + +print(f"\n✓ Saved: {OUTPUT_PREFIX}.pdf") +print(f"✓ Saved: {OUTPUT_PREFIX}.png") diff --git a/tools/build_fig1_separate.py b/tools/build_fig1_separate.py new file mode 100644 index 0000000000000000000000000000000000000000..38adbf2f58c95e4e7c83056a08c17870c30e2e6e --- /dev/null +++ b/tools/build_fig1_separate.py @@ -0,0 +1,191 @@ +import matplotlib.pyplot as plt +from PIL import Image, ImageDraw, ImageFont +import os +import glob +import json +import numpy as np +from pdf2image import convert_from_path + +# ========================================================================= +# 配置项 +# ========================================================================= +SCENE = "bonsai" +VIEW = "00000.png" +OUTPUTS_BASE = "/root/autodl-tmp/SplatAtlas/outputs" +METHODS = ["vanilla_3dgs", "pgsr", "ges", "opti3dgs", "steepgs"] + +LEFT_PDF_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf" +OUTPUT_DIR = "/root/autodl-tmp/SplatAtlas/tex/figures" + +# ========================================================================= +# 辅助函数 (构建右图 6宫格 用) +# ========================================================================= +def get_valid_dir(method, scene): + for suffix in ["", "_bak"]: + d = f"{OUTPUTS_BASE}/{method}_{scene}{suffix}" + if os.path.exists(d) and len(glob.glob(f"{d}/renders_test*")) > 0: + return d + return None + +def get_metrics(cell_dir): + try: + metric_files = sorted(glob.glob(f"{cell_dir}/metrics_test_iter*.json")) + if not metric_files: return None + with open(metric_files[-1], 'r') as f: + data = json.load(f) + geo = data.get("geometric", {}) + return { + "opacity": geo.get("alpha_mean", 0.0), + "aniso": geo.get("gamma_median", 0.0), + "psnr": data.get("photometric", {}).get("PSNR", 0.0) + } + except Exception: + return None + +def add_label(img, title, psnr, opacity=None, aniso=None, is_gt=False): + """在图片上绘制带有半透明黑底的白字标签""" + draw = ImageDraw.Draw(img, 'RGBA') + w, h = img.size + title_text = title.upper() + font_size = max(16, int(h * 0.04)) # 字号稍微调大一点点 + try: + font = ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) + except: + font = ImageFont.load_default() + + padding = 5 + title_bbox = draw.textbbox((0, 0), title_text, font=font) + tw, th = title_bbox[2] - title_bbox[0], title_bbox[3] - title_bbox[1] + title_x = (w - tw) // 2 + title_y = 10 + draw.rectangle([title_x - padding, title_y - padding, + title_x + tw + padding, title_y + th + padding], + fill=(0, 0, 0, 180)) + title_color = (255, 255, 255) if is_gt else (100, 255, 100) + draw.text((title_x, title_y), title_text, font=font, fill=title_color) + + if psnr: + psnr_text = f"PSNR: {psnr:.2f} dB" if not is_gt else "Reference" + p_bbox = draw.textbbox((0, 0), psnr_text, font=font) + pw, ph = p_bbox[2] - p_bbox[0], p_bbox[3] - p_bbox[1] + px = w - pw - 10 + py = h - ph - 45 # 留出底栏空间 + draw.rectangle([px - padding, py - padding, + px + pw + padding, py + ph + padding], + fill=(0, 0, 0, 180)) + draw.text((px, py), psnr_text, font=font, fill=(255,255,255)) + + if opacity is not None and aniso is not None: + bar_h = 35 + draw.rectangle([0, h - bar_h, w, h], fill=(30, 30, 30, 255)) + stat_text = f"opacity μ: {opacity:.3f} | anisotropy μ: {aniso:.3f}" if not is_gt else "opacity μ: 1.000 | anisotropy μ: 1.000" + s_bbox = draw.textbbox((0, 0), stat_text, font=font) + sw, sh = s_bbox[2] - s_bbox[0], s_bbox[3] - s_bbox[1] + draw.text(((w-sw)//2, h - bar_h + (bar_h-sh)//2), stat_text, font=font, fill=(200,200,200)) + + return img + +os.makedirs(OUTPUT_DIR, exist_ok=True) + +# ========================================================================= +# 处理 Fig 1a: 左图 (Rank Instability) 单独输出 +# ========================================================================= +print("正在处理 Fig 1a (左图)...") +if not os.path.exists(LEFT_PDF_PATH): + print(f"❌ 找不到左图 (PDF): {LEFT_PDF_PATH}") +else: + try: + pages = convert_from_path(LEFT_PDF_PATH, dpi=300) + img_left = pages[0] + + # 为左图单独创建一个更宽的画板 + fig_a, ax_a = plt.subplots(figsize=(12, 6)) + ax_a.imshow(img_left) + ax_a.axis('off') + ax_a.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", + fontsize=20, pad=15, loc='center', fontweight='bold') + + out_a_pdf = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.pdf") + out_a_png = os.path.join(OUTPUT_DIR, "fig1a_rank_instability.png") + + fig_a.savefig(out_a_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True) + fig_a.savefig(out_a_png, dpi=300, bbox_inches='tight', format='png', transparent=False) + print(f"✅ Fig 1a 保存成功: {out_a_pdf}") + plt.close(fig_a) + except Exception as e: + print(f"❌ 左图处理失败: {e}") + +# ========================================================================= +# 处理 Fig 1b: 右图 (6宫格 Render vs Rep) 单独输出 +# ========================================================================= +print("\n正在处理 Fig 1b (右图 6宫格)...") +images = [] +labels = ["GT"] + METHODS + +# 1. 提取 GT +gt_dir = get_valid_dir("vanilla_3dgs", SCENE) +if gt_dir: + gt_path = os.path.join(gt_dir, "gt_test_30000", VIEW) + if not os.path.exists(gt_path): + gt_path = os.path.join(gt_dir, "gt_test", VIEW) + if os.path.exists(gt_path): + img_gt = Image.open(gt_path).convert("RGBA") + img_gt = add_label(img_gt, "GT", None, opacity=1.0, aniso=1.0, is_gt=True) + images.append(img_gt) + else: + print(f"❌ 找不到 GT 图像: {gt_path}"); exit(1) +else: + print("❌ 找不到 vanilla_3dgs 目录以提取 GT"); exit(1) + +# 2. 提取 Methods +for method in METHODS: + cell_dir = get_valid_dir(method, SCENE) + if not cell_dir: + images.append(Image.new("RGBA", images[0].size, (50,50,50,255))) + continue + render_dirs = sorted(glob.glob(f"{cell_dir}/renders_test*")) + if not render_dirs: + images.append(Image.new("RGBA", images[0].size, (50,50,50,255))) + continue + img_path = os.path.join(render_dirs[-1], VIEW) + if not os.path.exists(img_path): + images.append(Image.new("RGBA", images[0].size, (50,50,50,255))) + continue + img_m = Image.open(img_path).convert("RGBA") + metrics = get_metrics(cell_dir) + if metrics: + img_m = add_label(img_m, method, metrics['psnr'], metrics['opacity'], metrics['aniso']) + else: + img_m = add_label(img_m, method, 0.0) + images.append(img_m) + +# 3. 拼接 6 宫格 +iw, ih = images[0].size +padding = 15 +grid_w = iw * 3 + padding * 2 +grid_h = ih * 2 + padding + +img_right = Image.new("RGBA", (grid_w, grid_h), (255,255,255,0)) # 使用透明底色 +positions = [ + (0, 0), (iw + padding, 0), ((iw + padding) * 2, 0), + (0, ih + padding), (iw + padding, ih + padding), ((iw + padding) * 2, ih + padding) +] +for img, pos in zip(images, positions): + img_right.paste(img, pos) + +# 4. 单独制图保存 +fig_b, ax_b = plt.subplots(figsize=(14, 8)) # 适合 3x2 的宽幅比例 +ax_b.imshow(img_right) +ax_b.axis('off') +ax_b.set_title(f"(b) Saturated Cluster ({SCENE.upper()}): Render vs. Representation", + fontsize=20, pad=15, loc='center', fontweight='bold') + +out_b_pdf = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.pdf") +out_b_png = os.path.join(OUTPUT_DIR, "fig1b_saturation_renders.png") + +fig_b.savefig(out_b_pdf, dpi=300, bbox_inches='tight', format='pdf', transparent=True) +fig_b.savefig(out_b_png, dpi=300, bbox_inches='tight', format='png', transparent=False) +print(f"✅ Fig 1b 保存成功: {out_b_pdf}") +plt.close(fig_b) + +print("\n=== 图片分离处理完成！===") diff --git a/tools/build_fig2.py b/tools/build_fig2.py new file mode 100644 index 0000000000000000000000000000000000000000..225f34900bc46b99f9132ed57adb54210dd9c7ad --- /dev/null +++ b/tools/build_fig2.py @@ -0,0 +1,107 @@ +import matplotlib as mpl +import matplotlib.pyplot as plt +import numpy as np +import os + +def create_variance_attribution_chart(): + # ========================================== + # 物理级环境重置：抹除所有全局样式污染 + # ========================================== + mpl.style.use('default') + mpl.rcParams.update(mpl.rcParamsDefault) + mpl.rcParams['axes.grid'] = False + mpl.rcParams['axes.axisbelow'] = True + mpl.rcParams['font.family'] = 'sans-serif' + mpl.rcParams['font.sans-serif'] = ['Helvetica', 'Arial'] + + # ========================================== + # 数据准备 + # ========================================== + categories = ['Opacity', 'Anisotropy', 'Scale', 'SH Energy Ratio', 'PSNR'] + method_variance = [86, 68, 43, 17, 1.4] + scene_variance = [10, 20, 45, 57, 98.6] + + categories.reverse() + method_variance.reverse() + scene_variance.reverse() + + color_method = '#4C72B0' + color_scene = '#DD8452' + + # ========================================== + # 画图配置 + # ========================================== + fig, ax = plt.subplots(figsize=(8, 4.5)) + + y_pos = np.arange(len(categories)) + bar_height = 0.6 + + bars_method = ax.barh(y_pos, method_variance, height=bar_height, + color=color_method, edgecolor='none', label=r'Method ($SS_{method}$)', zorder=3) + bars_scene = ax.barh(y_pos, scene_variance, left=method_variance, height=bar_height, + color=color_scene, edgecolor='none', label=r'Scene ($SS_{scene}$)', zorder=3) + + # ========================================== + # 细节修饰 + # ========================================== + ax.set_yticks(y_pos) + ax.set_yticklabels(categories, fontsize=11) + ax.set_xlim(0, 100) + + ax.set_xticks([0, 20, 40, 60, 80, 100]) + ax.set_xticklabels(['0%', '20%', '40%', '60%', '80%', '100%'], color='#555555') + + for i, (m_val, s_val) in enumerate(zip(method_variance, scene_variance)): + if m_val > 5: + ax.text(m_val / 2, i, f'{m_val}%' if m_val == int(m_val) else f'{m_val}%', + ha='center', va='center', color='white', fontweight='bold', fontsize=10, zorder=4) + elif m_val > 0: + ax.text(m_val + 2, i, f'{m_val}%', + ha='left', va='center', color=color_method, fontweight='bold', fontsize=10, zorder=4) + + if s_val > 5: + ax.text(m_val + s_val / 2, i, f'{s_val}%' if s_val == int(s_val) else f'{s_val}%', + ha='center', va='center', color='white', fontweight='bold', fontsize=10, zorder=4) + + # ========================================== + # 彻底的坐标轴与网格清理机制 (终极防御) + # ========================================== + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + ax.spines['left'].set_visible(False) + ax.spines['bottom'].set_linewidth(1.5) + ax.spines['bottom'].set_color('#333333') + + # 关闭 X 轴和 Y 轴的刻度短线 + ax.tick_params(axis='y', length=0) + ax.tick_params(axis='x', length=0) + + # 强制遍历底层实例并将其不可见化 + ax.grid(False) + for line in ax.get_xgridlines() + ax.get_ygridlines(): + line.set_visible(False) + line.set_color('none') + line.set_linewidth(0) + + # ========================================== + # 分割线 + # ========================================== + ax.axhline(y=0.5, color='#CCCCCC', linestyle='--', linewidth=1, zorder=2) + + ax.set_title("Asymmetric Variance Attribution", fontsize=14, fontweight='bold', pad=20) + ax.legend(loc='upper center', bbox_to_anchor=(0.5, 1.08), ncol=2, frameon=False, fontsize=11) + + plt.tight_layout() + + out_dir = "/root/autodl-tmp/SplatAtlas/tex/figures" + os.makedirs(out_dir, exist_ok=True) + out_pdf = os.path.join(out_dir, "fig2_variance_attribution.pdf") + out_png = os.path.join(out_dir, "fig2_variance_attribution.png") + + plt.savefig(out_pdf, format='pdf', bbox_inches='tight', dpi=300) + plt.savefig(out_png, format='png', bbox_inches='tight', dpi=300) + + print(f"✅ Fig 2 生成成功 (全局样式污染已重置，底层网格实例已物理擦除): {out_pdf}") + +if __name__ == "__main__": + create_variance_attribution_chart() diff --git a/tools/check_missing_renders.py b/tools/check_missing_renders.py new file mode 100644 index 0000000000000000000000000000000000000000..692af63d18bec1ba0d32553858d5c29478022e13 --- /dev/null +++ b/tools/check_missing_renders.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Diagnostic: check whether CoAdaptGS / LoD-GS have PLY, test renders, and metrics +for the cells that are missing from outputs/phase5a/task_vapre_splatatlas_repro.csv. +""" +import json +from pathlib import Path + +ROOT = Path("/root/autodl-tmp/SplatAtlas") +OUTPUTS = ROOT / "outputs" + +# Cells missing from phase5a (from previous diagnostic) +COADAPTGS_MISSING = [ + # T&T (12) + "Auditorium", "Ballroom", "Barn", "Caterpillar", "Courtroom", + "Lighthouse", "Museum", "Palace", "Playground", "Temple", "Train", "Truck", + # Mip-360 (9) + "Bicycle", "Bonsai", "Counter", "Flowers", "Garden", "Kitchen", + "Room", "Stump", "Treehill", + # DB (2) + "DrJohnson", "Playroom", +] + +LOD_GS_MISSING = [ + # NS (8) + "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship", +] + +# Common alternate casings to try, since phase5a uses lowercase but +# directory names sometimes use the canonical mixed case. +def candidate_dirs(method: str, scene: str): + cands = [ + OUTPUTS / f"{method}_{scene}", + OUTPUTS / f"{method}_{scene.lower()}", + OUTPUTS / f"{method}_{scene.capitalize()}", + ] + seen = [] + for c in cands: + if c not in seen: + seen.append(c) + return seen + +def find_run_dir(method: str, scene: str): + for d in candidate_dirs(method, scene): + if d.exists(): + return d + # Fallback: glob anything matching + matches = list(OUTPUTS.glob(f"{method}_*")) + for m in matches: + # match scene case-insensitively + suffix = m.name[len(method) + 1:] + if suffix.lower() == scene.lower(): + return m + return None + +def count_images(d: Path): + if not d.exists(): + return 0 + n = 0 + for ext in ("*.png", "*.jpg", "*.jpeg"): + n += len(list(d.glob(ext))) + return n + +def inspect_metrics_json(p: Path): + if not p.exists(): + return None + try: + obj = json.load(open(p, "r")) + except Exception as e: + return {"_error": str(e)} + flat = {} + def walk(o, prefix=""): + if isinstance(o, dict): + for k, v in o.items(): + walk(v, f"{prefix}{k}.") + elif isinstance(o, list): + pass + else: + flat[prefix.rstrip(".")] = o + walk(obj) + keep = {} + for k, v in flat.items(): + kl = k.lower() + if any(t in kl for t in ["psnr", "ssim", "lpips"]): + keep[k] = v + return keep or {"_no_metric_keys": list(flat.keys())[:10]} + +def check(method: str, scene: str): + run = find_run_dir(method, scene) + row = { + "method": method, + "scene": scene, + "run_dir": str(run) if run else "MISSING", + } + if run is None: + row.update({"ply": "-", "cameras": "-", "renders": "-", "gt": "-", + "metrics_json": "-", "metrics_values": "-"}) + return row + + ply = run / "point_cloud" / "iteration_30000" / "point_cloud.ply" + cameras = run / "cameras.json" + + # test render dirs (typical 3DGS-family layout) + test_root_candidates = [ + run / "test" / "ours_30000", + run / "test" / "ours_30000" / "30000", + ] + test_root = next((p for p in test_root_candidates if p.exists()), None) + + if test_root: + renders_dir = test_root / "renders" + gt_dir = test_root / "gt" + else: + # Fall back: any folder named renders / gt under run/ + renders_list = list(run.rglob("renders")) + gt_list = list(run.rglob("gt")) + renders_dir = renders_list[0] if renders_list else None + gt_dir = gt_list[0] if gt_list else None + + n_renders = count_images(renders_dir) if renders_dir else 0 + n_gt = count_images(gt_dir) if gt_dir else 0 + + metrics_path = run / "metrics_test_iter30000.json" + metrics_values = inspect_metrics_json(metrics_path) + + row.update({ + "ply": "OK" if ply.exists() else "MISSING", + "cameras": "OK" if cameras.exists() else "MISSING", + "renders": f"{n_renders} imgs @ {renders_dir}" if renders_dir else "MISSING", + "gt": f"{n_gt} imgs @ {gt_dir}" if gt_dir else "MISSING", + "metrics_json": "OK" if metrics_path.exists() else "MISSING", + "metrics_values": metrics_values if metrics_values else "-", + }) + return row + +def print_block(method: str, scenes): + print(f"\n{'='*100}") + print(f" {method} ({len(scenes)} missing scenes)") + print(f"{'='*100}") + summary = {"run_dir": 0, "ply": 0, "renders": 0, "metrics_json": 0} + for s in scenes: + r = check(method, s) + print(f"\n--- {method} × {s} ---") + print(f" run_dir : {r['run_dir']}") + print(f" PLY : {r['ply']}") + print(f" cameras.json : {r['cameras']}") + print(f" test renders : {r['renders']}") + print(f" test gt : {r['gt']}") + print(f" metrics_json : {r['metrics_json']}") + if isinstance(r['metrics_values'], dict): + print(f" metrics_values: {r['metrics_values']}") + if r['run_dir'] != "MISSING": + summary["run_dir"] += 1 + if r['ply'] == "OK": + summary["ply"] += 1 + if isinstance(r['renders'], str) and r['renders'].startswith(tuple("0123456789")): + n = int(r['renders'].split()[0]) + if n > 0: + summary["renders"] += 1 + if r['metrics_json'] == "OK": + summary["metrics_json"] += 1 + print(f"\n[Summary {method}]") + for k, v in summary.items(): + print(f" {k:<14}: {v} / {len(scenes)}") + +print_block("coadaptgs", COADAPTGS_MISSING) +print_block("lod_gs", LOD_GS_MISSING) diff --git a/tools/check_ply_and_find_render.py b/tools/check_ply_and_find_render.py new file mode 100644 index 0000000000000000000000000000000000000000..6586d810abd463b3dcc713a56c5dc596e9165cee --- /dev/null +++ b/tools/check_ply_and_find_render.py @@ -0,0 +1,197 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Two jobs: + 1. Verify the PLYs of the 31 missing-eval cells are real and readable. + 2. Locate render / eval scripts in the project and a reference run + that already has test/ours_30000/renders + gt as a template. +""" +import os +import re +import subprocess +from pathlib import Path + +ROOT = Path("/root/autodl-tmp/SplatAtlas") +OUTPUTS = ROOT / "outputs" + +CELLS = [] +for s in [ + "Auditorium","Ballroom","Barn","Caterpillar","Courtroom","Lighthouse", + "Museum","Palace","Playground","Temple","Train","Truck", + "Bicycle","Bonsai","Counter","Flowers","Garden","Kitchen", + "Room","Stump","Treehill","DrJohnson","Playroom", +]: + CELLS.append(("coadaptgs", s)) +for s in ["Chair","Drums","Ficus","Hotdog","Lego","Materials","Mic","Ship"]: + CELLS.append(("lod_gs", s)) + +def find_run(method, scene): + for cand in [ + OUTPUTS / f"{method}_{scene}", + OUTPUTS / f"{method}_{scene.lower()}", + ]: + if cand.exists(): + return cand + return None + +def parse_ply_header(p: Path): + """Return (n_vertices, n_props, format) by reading just the header.""" + try: + with open(p, "rb") as f: + magic = f.readline().strip() + if magic != b"ply": + return None + fmt = None + n_vertex = None + n_props = 0 + in_vertex = False + for _ in range(200): + line = f.readline() + if not line: + break + ls = line.decode("utf-8", errors="ignore").strip() + if ls.startswith("format "): + fmt = ls.split()[1] + if ls.startswith("element vertex"): + n_vertex = int(ls.split()[-1]) + in_vertex = True + elif ls.startswith("element ") and not ls.startswith("element vertex"): + in_vertex = False + elif in_vertex and ls.startswith("property"): + n_props += 1 + if ls == "end_header": + break + return (n_vertex, n_props, fmt) + except Exception as e: + return None + +# ============================================================ +# 1. Verify PLYs of the 31 missing cells +# ============================================================ +print("="*100) +print(" PLY integrity check for the 31 missing-eval cells") +print("="*100) +print(f"{'method':14s} {'scene':14s} {'size(MB)':>10s} {'vertices':>10s} {'props':>6s} {'fmt':>30s}") +bad = [] +for method, scene in CELLS: + run = find_run(method, scene) + if run is None: + print(f"{method:14s} {scene:14s} RUN DIR MISSING") + bad.append((method, scene, "no run dir")) + continue + ply = run / "point_cloud" / "iteration_30000" / "point_cloud.ply" + if not ply.exists(): + print(f"{method:14s} {scene:14s} PLY MISSING") + bad.append((method, scene, "no ply")) + continue + size_mb = ply.stat().st_size / 1024 / 1024 + info = parse_ply_header(ply) + if info is None: + print(f"{method:14s} {scene:14s} {size_mb:10.2f} HEADER UNREADABLE") + bad.append((method, scene, "header unreadable")) + continue + n_v, n_p, fmt = info + flag = "" + if n_v is None or n_v < 100: + flag = " <-- TOO FEW VERTICES" + bad.append((method, scene, f"only {n_v} vertices")) + print(f"{method:14s} {scene:14s} {size_mb:10.2f} {str(n_v):>10s} {n_p:>6d} {str(fmt):>30s}{flag}") + +print() +print(f"Bad PLYs: {len(bad)} / {len(CELLS)}") +for b in bad: + print(" ", b) + +# ============================================================ +# 2. Find render / eval scripts in the project +# ============================================================ +print() +print("="*100) +print(" Render / eval scripts in the repo") +print("="*100) + +candidates = [] +for p in ROOT.rglob("*.py"): + sp = str(p) + if "/.cache/" in sp or "/__pycache__/" in sp: + continue + name = p.name.lower() + if name in {"render.py", "render_test.py", "metrics.py", "evaluate.py", "eval.py"}: + candidates.append(p) + elif name.startswith("render") or name.startswith("metric") or name.startswith("eval"): + candidates.append(p) + +# Prefer ones at repo root or under known method dirs / scripts/ +def rank(p: Path): + sp = str(p) + score = 0 + if "/scripts/" in sp: score += 5 + if sp.endswith("/render.py"): score += 10 + if sp.endswith("/metrics.py"): score += 8 + if sp.endswith("/full_eval.py"): score += 6 + return -score + +candidates = sorted(set(candidates), key=rank) +for p in candidates[:50]: + print(" ", p) + +# ============================================================ +# 3. Find a reference run that already has test renders +# ============================================================ +print() +print("="*100) +print(" Reference runs that already have test/ours_30000/renders + gt") +print("="*100) + +found_refs = [] +for run in OUTPUTS.iterdir(): + if not run.is_dir(): + continue + if run.name.startswith("."): + continue + test = run / "test" / "ours_30000" + if not test.exists(): + continue + rdir = test / "renders" + gdir = test / "gt" + n_r = len(list(rdir.glob("*.png"))) if rdir.exists() else 0 + n_g = len(list(gdir.glob("*.png"))) if gdir.exists() else 0 + if n_r > 0 and n_g > 0: + found_refs.append((run.name, n_r, n_g)) + +print(f"Total reference runs with renders: {len(found_refs)}") +for name, nr, ng in found_refs[:20]: + print(f" {name:40s} renders={nr:4d} gt={ng:4d}") + +# Pick a same-scene reference for one of our missing cells if possible. +print() +print("Same-scene reference matches for missing cells (if any):") +ref_by_scene = {} +for name, nr, ng in found_refs: + # split first '_' is wrong because methods can have '_'. + # Try matching against canonical scenes used in this project. + pass + +CANON_SCENES = [ + "Auditorium","Ballroom","Barn","Caterpillar","Courtroom","Lighthouse", + "Museum","Palace","Playground","Temple","Train","Truck", + "Bicycle","Bonsai","Counter","Flowers","Garden","Kitchen","Room", + "Stump","Treehill","Chair","Drums","Ficus","Hotdog","Lego", + "Materials","Mic","Ship","DrJohnson","Playroom", +] +# Build {scene -> [example run names that have renders]} +ref_by_scene = {s: [] for s in CANON_SCENES} +for name, nr, ng in found_refs: + for s in CANON_SCENES: + suf = "_" + s + suf2 = "_" + s.lower() + if name.endswith(suf) or name.endswith(suf2): + ref_by_scene[s].append(name) + break + +for method, scene in CELLS: + refs = ref_by_scene.get(scene, []) + if refs: + print(f" need: {method:14s} × {scene:12s} -> reference: {refs[0]}") + else: + print(f" need: {method:14s} × {scene:12s} -> NO REFERENCE FOUND") diff --git a/tools/compose_fig1.py b/tools/compose_fig1.py new file mode 100644 index 0000000000000000000000000000000000000000..f6800fe01b7f9332f4da676d07d5c940aeb14804 --- /dev/null +++ b/tools/compose_fig1.py @@ -0,0 +1,72 @@ +import matplotlib.pyplot as plt +from PIL import Image +import os + +def compose_figure_1(left_path, right_path, output_prefix): + if not os.path.exists(left_path): + print(f"❌ 找不到左图: {left_path} \n👉 请手动修改脚本底部的 LEFT_PATH") + return + if not os.path.exists(right_path): + print(f"❌ 找不到右图: {right_path} \n👉 请手动修改脚本底部的 RIGHT_PATH") + return + + # 读取图像 + img_l = Image.open(left_path) + img_r = Image.open(right_path) + + # 统一高度对齐 (以 1000px 为基准高度进行 Lanczos 高质量重采样) + target_height = 1000 + aspect_l = img_l.width / img_l.height + aspect_r = img_r.width / img_r.height + + new_w_l = int(target_height * aspect_l) + new_w_r = int(target_height * aspect_r) + + img_l = img_l.resize((new_w_l, target_height), Image.Resampling.LANCZOS) + img_r = img_r.resize((new_w_r, target_height), Image.Resampling.LANCZOS) + + # 设置画布，保持两个 Panel 的原始宽高比比例 + fig, (ax1, ax2) = plt.subplots( + 1, 2, + figsize=(16, 5), + gridspec_kw={'width_ratios': [aspect_l, aspect_r]} + ) + + # 绘制 Left Panel + ax1.imshow(img_l) + ax1.axis('off') + ax1.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", + fontsize=16, pad=12, loc='left', fontweight='bold') + + # 绘制 Right Panel + ax2.imshow(img_r) + ax2.axis('off') + ax2.set_title("(b) Saturated Cluster: Render vs. Representation", + fontsize=16, pad=12, loc='left', fontweight='bold') + + # 缩小子图间距 + plt.subplots_adjust(wspace=0.03) + + # 导出保存 + os.makedirs(os.path.dirname(output_prefix), exist_ok=True) + pdf_path = f"{output_prefix}.pdf" + png_path = f"{output_prefix}.png" + + plt.savefig(pdf_path, dpi=300, bbox_inches='tight', format='pdf', transparent=True) + plt.savefig(png_path, dpi=300, bbox_inches='tight', format='png', transparent=False) + + print(f"=== Figure 1 (Saturation Puzzle) 拼图合成完毕 ===") + print(f"✅ Saved PDF: {pdf_path} (供 LaTeX 插入)") + print(f"✅ Saved PNG: {png_path} (供快速预览)") + +if __name__ == "__main__": + # ========================================================================= + # ⚠️ 请在这里替换为您实际的文件路径 + # ========================================================================= + LEFT_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.png" + RIGHT_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/image_1ae31e.jpg" + + # 自动保存到论文的 tex figures 目录 + OUTPUT_PREFIX = "/root/autodl-tmp/SplatAtlas/tex/figures/fig1_saturation_puzzle" + + compose_figure_1(LEFT_PATH, RIGHT_PATH, OUTPUT_PREFIX) diff --git a/tools/compose_fig1_pdf.py b/tools/compose_fig1_pdf.py new file mode 100644 index 0000000000000000000000000000000000000000..ad5f48e94d68e526d0388f50ea4982133b34bc54 --- /dev/null +++ b/tools/compose_fig1_pdf.py @@ -0,0 +1,79 @@ +import matplotlib.pyplot as plt +from PIL import Image +import os +from pdf2image import convert_from_path + +def compose_figure_1(left_pdf_path, right_img_path, output_prefix): + if not os.path.exists(left_pdf_path): + print(f"❌ 找不到左图 (PDF): {left_pdf_path}") + return + if not os.path.exists(right_img_path): + print(f"❌ 找不到右图 (IMG): {right_img_path}") + return + + print("正在从 PDF 提取左侧面板 (DPI=300)...") + # 将 PDF 的第一页转换为 PIL Image + pages = convert_from_path(left_pdf_path, dpi=300) + img_l = pages[0] + + print("正在加载右侧面板...") + img_r = Image.open(right_img_path) + + # 统一高度对齐 (以 1500px 为基准高度进行高质量重采样，保证 PDF 提取的清晰度) + target_height = 1500 + aspect_l = img_l.width / img_l.height + aspect_r = img_r.width / img_r.height + + new_w_l = int(target_height * aspect_l) + new_w_r = int(target_height * aspect_r) + + img_l = img_l.resize((new_w_l, target_height), Image.Resampling.LANCZOS) + img_r = img_r.resize((new_w_r, target_height), Image.Resampling.LANCZOS) + + # 设置画布，保持两个 Panel 的原始宽高比比例 + # 使用更大的 figsize 以容纳高分辨率图像 + fig, (ax1, ax2) = plt.subplots( + 1, 2, + figsize=(18, 6), + gridspec_kw={'width_ratios': [aspect_l, aspect_r]} + ) + + # 绘制 Left Panel + ax1.imshow(img_l) + ax1.axis('off') + ax1.set_title("(a) Per-seed PSNR Rank Instability (BONSAI & LEGO)", + fontsize=18, pad=12, loc='left', fontweight='bold') + + # 绘制 Right Panel + ax2.imshow(img_r) + ax2.axis('off') + ax2.set_title("(b) Saturated Cluster: Render vs. Representation", + fontsize=18, pad=12, loc='left', fontweight='bold') + + # 缩小子图间距 + plt.subplots_adjust(wspace=0.03) + + # 导出保存 + os.makedirs(os.path.dirname(output_prefix), exist_ok=True) + pdf_path = f"{output_prefix}.pdf" + png_path = f"{output_prefix}.png" + + print("正在保存最终拼接文件...") + plt.savefig(pdf_path, dpi=300, bbox_inches='tight', format='pdf', transparent=True) + plt.savefig(png_path, dpi=300, bbox_inches='tight', format='png', transparent=False) + + print(f"=== Figure 1 (Saturation Puzzle) 拼图合成完毕 ===") + print(f"✅ Saved PDF: {pdf_path} (供 LaTeX 插入)") + print(f"✅ Saved PNG: {png_path} (供快速预览)") + +if __name__ == "__main__": + # ========================================================================= + # 文件路径配置 + # ========================================================================= + LEFT_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/fig1_left_panel.pdf" + RIGHT_PATH = "/root/autodl-tmp/SplatAtlas/outputs/phase5b/image_1ae31e.jpg" + + # 自动保存到论文的 tex figures 目录 + OUTPUT_PREFIX = "/root/autodl-tmp/SplatAtlas/tex/figures/fig1_saturation_puzzle" + + compose_figure_1(LEFT_PATH, RIGHT_PATH, OUTPUT_PREFIX) diff --git a/tools/convert_phase5a_to_appE_long.py b/tools/convert_phase5a_to_appE_long.py new file mode 100644 index 0000000000000000000000000000000000000000..25746fb1e5a22cad14a77607deae5423354dca5b --- /dev/null +++ b/tools/convert_phase5a_to_appE_long.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import csv +import math +from collections import defaultdict +from pathlib import Path + +IN = Path("outputs/phase5a/task_vapre_splatatlas_repro.csv") +OUT = Path("outputs/phase1/per_cell_metrics_long.csv") +OUT.parent.mkdir(parents=True, exist_ok=True) + +METHOD_COL = "method" +SCENE_COL = "scene" + +METRICS = { + "psnr_splat": "psnr", + "ssim_splat": "ssim", + "lpips_splat": "lpips", +} + +EXPECTED_METHODS = [ + "3dgsmcmc", "absgs", "atomgs", "conegs", "ges", "ghap", "gslpm", + "lapisgs", "reactgs", "vanilla_3dgs", + "2dgs", "gaussian_surfel", "gof", "pgsr", "scaffoldgs", + "coadaptgs", "erankgs", "gaussianpro", "minisplatting", + "opti3dgs", "steepgs", + "3dgs_dr", "analyticsplatting", "lod_gs", "mipsplatting", + "pixelgs", + "cdcgs", "hogs", "lightgaussian", "octree_gs", "trimgs", +] + +EXPECTED_SCENES = [ + "auditorium", "ballroom", "barn", "caterpillar", "courtroom", + "lighthouse", "museum", "palace", "playground", "temple", "train", "truck", + "bicycle", "bonsai", "counter", "flowers", "garden", "kitchen", "room", + "stump", "treehill", + "chair", "drums", "ficus", "hotdog", "lego", "materials", "mic", "ship", + "drjohnson", "playroom", +] + +def parse_float(x): + x = str(x).strip() + if x == "" or x.lower() in {"nan", "none", "null", "na", "n/a", "---"}: + return None + try: + v = float(x) + except ValueError: + return None + if math.isnan(v) or math.isinf(v): + return None + return v + +buckets = defaultdict(list) + +with open(IN, "r", newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + fields = set(reader.fieldnames or []) + + required = {METHOD_COL, SCENE_COL, *METRICS.keys()} + missing = required - fields + if missing: + raise ValueError(f"Missing columns in {IN}: {missing}. Actual columns={reader.fieldnames}") + + for row in reader: + method = row[METHOD_COL].strip() + scene = row[SCENE_COL].strip() + + for src_col, metric in METRICS.items(): + value = parse_float(row[src_col]) + if value is None: + continue + buckets[(method, scene, metric)].append(value) + +rows = [] +for (method, scene, metric), vals in sorted(buckets.items()): + # If duplicate / multi-seed rows exist, average them. + value = sum(vals) / len(vals) + rows.append({ + "method_key": method, + "scene": scene, + "metric": metric, + "value": f"{value:.10g}", + }) + +with open(OUT, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=["method_key", "scene", "metric", "value"]) + writer.writeheader() + writer.writerows(rows) + +methods = sorted({r["method_key"] for r in rows}) +scenes = sorted({r["scene"] for r in rows}) + +print("Wrote:", OUT) +print("rows:", len(rows)) +print("methods:", len(methods), methods) +print("scenes:", len(scenes), scenes) + +for metric in ["psnr", "ssim", "lpips"]: + cells = {(r["method_key"], r["scene"]) for r in rows if r["metric"] == metric} + print(metric, "cells:", len(cells)) + +missing_methods = sorted(set(EXPECTED_METHODS) - set(methods)) +extra_methods = sorted(set(methods) - set(EXPECTED_METHODS)) +missing_scenes = sorted(set(EXPECTED_SCENES) - set(scenes)) +extra_scenes = sorted(set(scenes) - set(EXPECTED_SCENES)) + +print("missing_methods:", missing_methods) +print("extra_methods:", extra_methods) +print("missing_scenes:", missing_scenes) +print("extra_scenes:", extra_scenes) + +print("erankgs × lego:") +for r in rows: + if r["method_key"] == "erankgs" and r["scene"].lower() == "lego": + print(r) + +if len(methods) != 31: + raise AssertionError(f"Expected 31 methods, got {len(methods)}") +if len(scenes) != 31: + raise AssertionError(f"Expected 31 scenes, got {len(scenes)}") +if missing_methods or extra_methods: + raise AssertionError(f"Method mismatch: missing={missing_methods}, extra={extra_methods}") +if missing_scenes or extra_scenes: + raise AssertionError(f"Scene mismatch: missing={missing_scenes}, extra={extra_scenes}") + +expected_rows = 31 * 31 * 3 +if len(rows) != expected_rows: + print(f"[WARN] Expected {expected_rows} metric rows, got {len(rows)}.") + print("[WARN] Missing metric values will be skipped in dataset means.") + print("[WARN] If an entire method-dataset block is missing, it will appear as --- in Appendix E.") diff --git a/tools/eval_all_31.sh b/tools/eval_all_31.sh new file mode 100644 index 0000000000000000000000000000000000000000..9dc017a70164c0dc85ec2f6ca9e6f27096680970 --- /dev/null +++ b/tools/eval_all_31.sh @@ -0,0 +1,125 @@ +#!/bin/bash +source /root/miniconda3/etc/profile.d/conda.sh +ROOT=/root/autodl-tmp/SplatAtlas +ITER=30000 + +# 环境及源码路径 +declare -A ENV_MAP CODE_MAP +ENV_MAP[coadaptgs]="/root/autodl-tmp/envs/co_adaptation_3dgs" +CODE_MAP[coadaptgs]="/root/autodl-tmp/CoAdaptGS" +ENV_MAP[lod_gs]="/root/autodl-tmp/envs/lod_gs" +CODE_MAP[lod_gs]="/root/autodl-tmp/LOD_GS" + +# 场景→数据集映射（仅需用来拼 colmap_dir，eval 不需要 source_path） +get_dataset_path() { + case "$1" in + auditorium|ballroom|barn|caterpillar|courtroom|lighthouse|museum|palace|playground|temple|train|truck) + echo "/root/autodl-tmp/dataset/tnt/$1" ;; + bicycle|flowers|garden|stump|treehill) + echo "/root/autodl-tmp/dataset/360/$1" ;; + bonsai|counter|kitchen|room) + echo "/root/autodl-tmp/dataset/360/$1" ;; + Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship) + echo "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/$1" ;; + DrJohnson|Playroom) + echo "/root/autodl-tmp/dataset/deepblending_clean/$1" ;; + *) echo "" ;; + esac +} + +# 全部 31 个 cell (method scene) +CELLS=( + "coadaptgs auditorium" + "coadaptgs ballroom" + "coadaptgs barn" + "coadaptgs caterpillar" + "coadaptgs courtroom" + "coadaptgs lighthouse" + "coadaptgs museum" + "coadaptgs palace" + "coadaptgs playground" + "coadaptgs temple" + "coadaptgs train" + "coadaptgs truck" + "coadaptgs bicycle" + "coadaptgs bonsai" + "coadaptgs counter" + "coadaptgs flowers" + "coadaptgs garden" + "coadaptgs kitchen" + "coadaptgs room" + "coadaptgs stump" + "coadaptgs treehill" + "coadaptgs DrJohnson" + "coadaptgs Playroom" + "lod_gs Chair" + "lod_gs Drums" + "lod_gs Ficus" + "lod_gs Hotdog" + "lod_gs Lego" + "lod_gs Materials" + "lod_gs Mic" + "lod_gs Ship" +) + +printf "%-18s %-20s %8s %8s %8s\n" "Method" "Scene" "PSNR" "SSIM" "LPIPS" +echo "--------------------------------------------------------------------" + +for entry in "${CELLS[@]}"; do + read -r method scene <<< "$entry" + env_path="${ENV_MAP[$method]}" + code_path="${CODE_MAP[$method]}" + if [ -z "$env_path" ] || [ -z "$code_path" ]; then + echo "Skipping $method/$scene (no env/code)" + continue + fi + + model_path="$ROOT/outputs/${method}_${scene}" + renders_dir="$model_path/renders_test_${ITER}" + gt_dir="$model_path/gt_test_${ITER}" + ply_path="$model_path/point_cloud/iteration_${ITER}/point_cloud.ply" + metrics_json="$model_path/metrics_test_iter${ITER}.json" + colmap_dir=$(get_dataset_path "$scene") + + # 跳过完全不存在渲染的 + if [ ! -d "$renders_dir" ] || [ $(ls "$renders_dir"/*.png 2>/dev/null | wc -l) -eq 0 ]; then + printf "%-18s %-20s %8s %8s %8s\n" "$method" "$scene" "NO_RENDER" "" "" + continue + fi + + # 激活环境 + conda activate "$env_path" + export PYTHONPATH="${code_path}:${ROOT}:${PYTHONPATH}" + + # 若 metrics json 不存在则计算 + if [ ! -f "$metrics_json" ]; then + python "$ROOT/ufd_evalkit/run_eval.py" \ + --method "$method" \ + --scene "$scene" \ + --render_dir "$renders_dir" \ + --gt_dir "$gt_dir" \ + --ply_path "$ply_path" \ + --output_json "$metrics_json" \ + --colmap_dir "$colmap_dir" 2>/dev/null + fi + + # 解析并打印 + if [ -f "$metrics_json" ]; then + read -r psnr ssim lpips < <( + python -c " +import json +with open('$metrics_json') as f: + d = json.load(f) +psnr = d['photometric'].get('PSNR', float('nan')) +ssim = d['photometric'].get('SSIM', float('nan')) +lpips = d['photometric'].get('LPIPS', float('nan')) +print(f'{psnr:.4f} {ssim:.4f} {lpips:.4f}') +" + ) + printf "%-18s %-20s %8s %8s %8s\n" "$method" "$scene" "$psnr" "$ssim" "$lpips" + else + printf "%-18s %-20s %8s %8s %8s\n" "$method" "$scene" "FAILED" "" "" + fi + + conda deactivate +done diff --git a/tools/fix_all_missing_renders.sh b/tools/fix_all_missing_renders.sh new file mode 100644 index 0000000000000000000000000000000000000000..b3013d8d23e335aa5844960ec74f00d1e7d7ebba --- /dev/null +++ b/tools/fix_all_missing_renders.sh @@ -0,0 +1,148 @@ +#!/bin/bash +set -e +source /root/miniconda3/etc/profile.d/conda.sh + +ROOT=/root/autodl-tmp/SplatAtlas +ITER=30000 + +# 方法 -> 环境 + 源码根目录 +declare -A ENV_MAP +declare -A CODE_MAP + +ENV_MAP[lod_gs]="/root/autodl-tmp/envs/lod_gs" +CODE_MAP[lod_gs]="/root/autodl-tmp/LOD_GS" + +ENV_MAP[coadaptgs]="/root/autodl-tmp/envs/co_adaptation_3dgs" +CODE_MAP[coadaptgs]="/root/autodl-tmp/CoAdaptGS" + +get_source_res() { + case "$1" in + auditorium|ballroom|barn|caterpillar|courtroom|lighthouse|museum|palace|playground|temple|train|truck) + echo "/root/autodl-tmp/dataset/tnt/$1 2" ;; + bicycle|flowers|garden|stump|treehill) + echo "/root/autodl-tmp/dataset/360/$1 4" ;; + bonsai|counter|kitchen|room) + echo "/root/autodl-tmp/dataset/360/$1 2" ;; + Chair|Drums|Ficus|Hotdog|Lego|Materials|Mic|Ship) + echo "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF/$1 1" ;; + DrJohnson|Playroom) + echo "/root/autodl-tmp/dataset/deepblending_clean/$1 1" ;; + *) + echo ""; return 1 ;; + esac +} + +MISSING=( + "coadaptgs auditorium" + "coadaptgs ballroom" + "coadaptgs barn" + "coadaptgs caterpillar" + "coadaptgs courtroom" + "coadaptgs lighthouse" + "coadaptgs museum" + "coadaptgs palace" + "coadaptgs playground" + "coadaptgs temple" + "coadaptgs train" + "coadaptgs truck" + "coadaptgs bicycle" + "coadaptgs bonsai" + "coadaptgs counter" + "coadaptgs flowers" + "coadaptgs garden" + "coadaptgs kitchen" + "coadaptgs room" + "coadaptgs stump" + "coadaptgs treehill" + "coadaptgs DrJohnson" + "coadaptgs Playroom" + "lod_gs Chair" + "lod_gs Drums" + "lod_gs Ficus" + "lod_gs Hotdog" + "lod_gs Lego" + "lod_gs Materials" + "lod_gs Mic" + "lod_gs Ship" +) + +run_step() { + local name="$1" + local cmd="$2" + echo " [Step] $name ..." + if eval "$cmd"; then + echo " [OK]" + return 0 + else + echo " [FAIL]" + return 1 + fi +} + +for entry in "${MISSING[@]}"; do + read -r method scene <<< "$entry" + env_path="${ENV_MAP[$method]}" + code_root="${CODE_MAP[$method]}" + if [ -z "$env_path" ] || [ -z "$code_root" ]; then + echo "Unknown method or code root for $method"; continue + fi + + read -r source_path resolution < <(get_source_res "$scene") + if [ -z "$source_path" ]; then + echo "Skipping $method/$scene: unknown dataset"; continue + fi + + model_path="$ROOT/outputs/${method}_${scene}" + ply_path="$model_path/point_cloud/iteration_${ITER}/point_cloud.ply" + if [ ! -f "$ply_path" ]; then + echo "Skipping $method/$scene: PLY not found"; continue + fi + + render_dir="$model_path/renders_test_${ITER}" + gt_dir="$model_path/gt_test_${ITER}" + render_flag="$model_path/render_complete_${ITER}.flag" + + png_count=$(find "$render_dir" -maxdepth 1 -name "*.png" 2>/dev/null | wc -l) + if [ "$png_count" -eq 0 ]; then + echo "=== $method / $scene (res=$resolution) - RENDER ===" + rm -rf "$render_dir" "$gt_dir" + else + echo "=== $method / $scene - renders exist ($png_count pngs) ===" + fi + + conda activate "$env_path" + + # 关键修复：把方法自己的源码根目录放在最前面 + export PYTHONPATH="${code_root}:${ROOT}:${PYTHONPATH:+:$PYTHONPATH}" + + if [ "$png_count" -eq 0 ]; then + run_step "Render" " + python $ROOT/scripts/main_render.py \ + --method $method \ + --source_path $source_path \ + --model_path $model_path \ + --iteration $ITER \ + --resolution $resolution + " || continue + fi + + metrics_json="$model_path/metrics_test_iter${ITER}.json" + if [ ! -f "$metrics_json" ]; then + run_step "Eval" " + python $ROOT/ufd_evalkit/run_eval.py \ + --method $method \ + --scene $scene \ + --render_dir $render_dir \ + --gt_dir $gt_dir \ + --ply_path $ply_path \ + --output_json $metrics_json \ + --colmap_dir $source_path + " || continue + else + echo " [Metrics json already exists]" + fi + + conda deactivate +done + +echo "All done. Now re-run fix_phase5a_table.py to update the CSV." diff --git a/tools/full_31x31_render_status.py b/tools/full_31x31_render_status.py new file mode 100644 index 0000000000000000000000000000000000000000..1f8aace17df40ca2a4026d26fe1f971acff251dc --- /dev/null +++ b/tools/full_31x31_render_status.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +import json +from pathlib import Path +from collections import defaultdict + +ROOT = Path("/root/autodl-tmp/SplatAtlas") +OUTPUTS = ROOT / "outputs" + +# 31 methods (完整列表) +METHODS = [ + "3dgsmcmc", "absgs", "atomgs", "conegs", "ges", "ghap", "gslpm", "lapisgs", "reactgs", "vanilla_3dgs", + "2dgs", "gaussian_surfel", "gof", "pgsr", "scaffoldgs", + "coadaptgs", "erankgs", "gaussianpro", "minisplatting", "opti3dgs", "steepgs", + "3dgs_dr", "analyticsplatting", "lod_gs", "mipsplatting", "pixelgs", + "cdcgs", "hogs", "lightgaussian", "octree_gs", "trimgs" +] + +# 31 canonical scenes +SCENES = [ + "Auditorium", "Ballroom", "Barn", "Caterpillar", "Courtroom", "Lighthouse", "Museum", "Palace", "Playground", "Temple", "Train", "Truck", + "Bicycle", "Bonsai", "Counter", "Flowers", "Garden", "Kitchen", "Room", "Stump", "Treehill", + "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship", + "DrJohnson", "Playroom" +] + +def check_cell(method, scene): + run_dir = OUTPUTS / f"{method}_{scene}" + if not run_dir.exists(): + return {"status": "NO_RUN_DIR"} + + renders_dir = run_dir / "renders_test_30000" + gt_dir = run_dir / "gt_test_30000" + metrics_file = run_dir / "metrics_test_iter30000.json" + + n_renders = len(list(renders_dir.glob("*.png"))) if renders_dir.exists() else 0 + n_gt = len(list(gt_dir.glob("*.png"))) if gt_dir.exists() else 0 + + metrics_status = "NO_METRICS" + psnr = ssim = lpips = None + if metrics_file.exists(): + try: + obj = json.load(open(metrics_file)) + photo = obj.get("photometric", {}) + psnr = photo.get("PSNR") + ssim = photo.get("SSIM") + lpips = photo.get("LPIPS") + + if psnr is not None and not (isinstance(psnr, float) and str(psnr) == "nan"): + metrics_status = "VALID" + else: + metrics_status = "NAN" + except: + metrics_status = "JSON_ERROR" + + return { + "status": "OK", + "n_renders": n_renders, + "n_gt": n_gt, + "metrics_status": metrics_status, + "psnr": psnr, + "ssim": ssim, + "lpips": lpips + } + +print("=" * 120) +print(f"{'method':14s} {'scene':12s} {'run?':5s} {'renders':8s} {'gt':6s} {'metrics':10s} {'PSNR':>8s} {'SSIM':>8s} {'LPIPS':>8s}") +print("=" * 120) + +summary = defaultdict(int) +for method in METHODS: + for scene in SCENES: + res = check_cell(method, scene) + summary[res["status"]] += 1 + if res["status"] == "OK": + summary[res["metrics_status"]] += 1 + + psnr_str = f"{res.get('psnr', ''):.2f}" if res.get("psnr") not in (None, "nan") else "" + ssim_str = f"{res.get('ssim', ''):.4f}" if res.get("ssim") not in (None, "nan") else "" + lpips_str = f"{res.get('lpips', ''):.4f}" if res.get("lpips") not in (None, "nan") else "" + + print(f"{method:14s} {scene:12s} {res['status']:5s} {res.get('n_renders', 0):>8d} {res.get('n_gt', 0):>6d} {res['metrics_status']:10s} {psnr_str:>8s} {ssim_str:>8s} {lpips_str:>8s}") + +print("\n" + "=" * 120) +print("SUMMARY") +print("=" * 120) +for k, v in sorted(summary.items()): + print(f"{k:20s}: {v}") +print(f"Total cells checked: {len(METHODS) * len(SCENES)}") diff --git a/tools/generate_appE_rows.py b/tools/generate_appE_rows.py new file mode 100644 index 0000000000000000000000000000000000000000..46f8f1e8c5f39d5c2a7090a82801097e924f46be --- /dev/null +++ b/tools/generate_appE_rows.py @@ -0,0 +1,370 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Generate Appendix E row-only LaTeX fragments for headline PSNR / SSIM / LPIPS. + +Expected input CSV schema, long format: + + method_key,scene,metric,value + 3dgsmcmc,Bonsai,psnr,32.398 + 3dgsmcmc,Bonsai,ssim,0.9421 + 3dgsmcmc,Bonsai,lpips,0.0584 + +If multiple rows exist for the same (method_key, scene, metric), they are averaged +first. This supports multi-seed cells by taking the seed mean. + +If your actual CSV is wide format or uses different column names, replace only +load_long_csv() below. Keep the aggregation / LaTeX logic unchanged. + +Usage: + python tools/generate_appE_rows.py path/to/per_cell_metrics_long.csv + +Outputs: + tex/appE_psnr_rows.tex + tex/appE_ssim_rows.tex + tex/appE_lpips_rows.tex +""" + +from __future__ import annotations + +import argparse +import csv +import math +import os +from collections import defaultdict +from typing import Dict, Iterable, List, Tuple, Optional + + +DEFAULT_INPUT_CSV = "outputs/phase1/per_cell_metrics_long.csv" +OUT_DIR = "tex" + +MASK_METHOD_FAILURE = False +METHOD_FAILURE_CELLS = { + ("erankgs", "Lego"), +} + +DATASETS = { + "T&T": [ + "Auditorium", "Ballroom", "Barn", "Caterpillar", "Courtroom", "Lighthouse", + "Museum", "Palace", "Playground", "Temple", "Train", "Truck", + ], + "Mip-360": [ + "Bicycle", "Bonsai", "Counter", "Flowers", "Garden", "Kitchen", "Room", + "Stump", "Treehill", + ], + "NS": [ + "Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship", + ], + "DB": [ + "DrJohnson", "Playroom", + ], +} + +DATASET_ORDER = ["T&T", "Mip-360", "NS", "DB"] +ALL_CANONICAL_SCENES = [scene for ds in DATASET_ORDER for scene in DATASETS[ds]] + +SCENE_CANONICAL = {s.lower(): s for s in ALL_CANONICAL_SCENES} +SCENE_CANONICAL.update({ + "auditorium": "Auditorium", + "ballroom": "Ballroom", + "barn": "Barn", + "caterpillar": "Caterpillar", + "courtroom": "Courtroom", + "lighthouse": "Lighthouse", + "museum": "Museum", + "palace": "Palace", + "playground": "Playground", + "temple": "Temple", + "train": "Train", + "truck": "Truck", + "bicycle": "Bicycle", + "bonsai": "Bonsai", + "counter": "Counter", + "flowers": "Flowers", + "garden": "Garden", + "kitchen": "Kitchen", + "room": "Room", + "stump": "Stump", + "treehill": "Treehill", + "chair": "Chair", + "drums": "Drums", + "ficus": "Ficus", + "hotdog": "Hotdog", + "lego": "Lego", + "materials": "Materials", + "mic": "Mic", + "ship": "Ship", + "drjohnson": "DrJohnson", + "dr_johnson": "DrJohnson", + "playroom": "Playroom", +}) + +CATEGORIES = [ + ("FREE", [ + "3dgsmcmc", "absgs", "atomgs", "conegs", "ges", "ghap", "gslpm", + "lapisgs", "reactgs", "vanilla_3dgs", + ]), + ("GEO_REG", [ + "2dgs", "gaussian_surfel", "gof", "pgsr", "scaffoldgs", + ]), + ("PRIM_CTRL", [ + "coadaptgs", "erankgs", "gaussianpro", "minisplatting", + "opti3dgs", "steepgs", + ]), + ("RENDER", [ + "3dgs_dr", "analyticsplatting", "lod_gs", "mipsplatting", + "pixelgs", + ]), + ("COMPRESS", [ + "cdcgs", "hogs", "lightgaussian", "octree_gs", "trimgs", + ]), +] + +DISPLAY_NAMES = { + "3dgsmcmc": "3DGS-MCMC", + "absgs": "AbsGS", + "atomgs": "AtomGS", + "conegs": "ConeGS", + "ges": "GES", + "ghap": "GHAP", + "gslpm": "GSLPM", + "lapisgs": "LapisGS", + "reactgs": "ReactGS", + "vanilla_3dgs": "3DGS", + "2dgs": "2DGS", + "gaussian_surfel": "Gaussian Surfels", + "gof": "GOF", + "pgsr": "PGSR", + "scaffoldgs": "Scaffold-GS", + "coadaptgs": "CoAdaptGS", + "erankgs": "eRankGS", + "gaussianpro": "GaussianPro", + "minisplatting": "MiniSplatting", + "opti3dgs": "Opti3DGS", + "steepgs": "SteepGS", + "3dgs_dr": "3DGS-DR", + "analyticsplatting": "AnalyticSplatting", + "lod_gs": "LoD-GS", + "mipsplatting": "Mip-Splatting", + "pixelgs": "PixelGS", + "cdcgs": "CDCGS", + "hogs": "HoGS", + "lightgaussian": "LightGaussian", + "octree_gs": "Octree-GS", + "trimgs": "TrimGS", +} + +METRIC_FORMATS = { + "psnr": "{:.2f}", + "ssim": "{:.4f}", + "lpips": "{:.4f}", +} + +OUTPUT_FILES = { + "psnr": "appE_psnr_rows.tex", + "ssim": "appE_ssim_rows.tex", + "lpips": "appE_lpips_rows.tex", +} + + +def latex_escape(text: str) -> str: + replacements = { + "\\": r"\textbackslash{}", + "_": r"\_", + "&": r"\&", + "%": r"\%", + "$": r"\$", + "#": r"\#", + "{": r"\{", + "}": r"\}", + } + return "".join(replacements.get(ch, ch) for ch in text) + + +def canonical_scene(scene: str) -> Optional[str]: + scene = scene.strip() + return SCENE_CANONICAL.get(scene.lower()) + + +def parse_float(value: str) -> Optional[float]: + value = value.strip() + if value == "" or value.lower() in {"nan", "none", "null", "na", "n/a", "---"}: + return None + try: + x = float(value) + except ValueError: + return None + if math.isnan(x) or math.isinf(x): + return None + return x + + +def mean_or_none(values: Iterable[float]) -> Optional[float]: + vals = list(values) + if not vals: + return None + return sum(vals) / len(vals) + + +def fmt_metric(metric: str, value: Optional[float]) -> str: + if value is None: + return "---" + return METRIC_FORMATS[metric].format(value) + + +def load_long_csv(path: str) -> Dict[Tuple[str, str, str], float]: + """ + Load long-format per-cell metrics and average duplicate rows. + + Returns: + dict[(method_key, canonical_scene, metric)] = mean_value + + TODO if your actual CSV is wide format: + Replace this function so it emits the same dictionary, e.g. + for columns method, scene, psnr, ssim, lpips: + out[(method, scene, "psnr")] = psnr_value + out[(method, scene, "ssim")] = ssim_value + out[(method, scene, "lpips")] = lpips_value + """ + buckets: Dict[Tuple[str, str, str], List[float]] = defaultdict(list) + + with open(path, "r", newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + if reader.fieldnames is None: + raise ValueError(f"Input CSV has no header: {path}") + + fields = set(reader.fieldnames) + method_col = "method_key" if "method_key" in fields else "method" + scene_col = "scene" + metric_col = "metric" + value_col = "value" + + required = {method_col, scene_col, metric_col, value_col} + missing = required - fields + if missing: + raise ValueError( + f"Input CSV missing columns {sorted(missing)}. " + f"Expected long schema: method_key,scene,metric,value. " + f"Actual columns: {reader.fieldnames}" + ) + + for row in reader: + method = row[method_col].strip() + scene_raw = row[scene_col].strip() + metric = row[metric_col].strip().lower() + value = parse_float(row[value_col]) + + if metric not in {"psnr", "ssim", "lpips"}: + continue + if value is None: + continue + + scene = canonical_scene(scene_raw) + if scene is None: + continue + + if MASK_METHOD_FAILURE and (method, scene) in METHOD_FAILURE_CELLS: + continue + + buckets[(method, scene, metric)].append(value) + + return {key: mean_or_none(vals) for key, vals in buckets.items() if vals} + + +def aggregate_for_method( + cell_values: Dict[Tuple[str, str, str], float], + method: str, + metric: str, +) -> Tuple[Optional[float], Optional[float], Optional[float], Optional[float], Optional[float]]: + dataset_means: List[Optional[float]] = [] + + for ds in DATASET_ORDER: + vals = [] + for scene in DATASETS[ds]: + v = cell_values.get((method, scene, metric)) + if v is not None: + vals.append(v) + dataset_means.append(mean_or_none(vals)) + + overall_vals = [] + for scene in ALL_CANONICAL_SCENES: + v = cell_values.get((method, scene, metric)) + if v is not None: + overall_vals.append(v) + + overall = mean_or_none(overall_vals) + return (*dataset_means, overall) + + +def render_metric_fragment( + cell_values: Dict[Tuple[str, str, str], float], + metric: str, +) -> List[str]: + lines: List[str] = [] + lines.append("% AUTOGENERATED by tools/generate_appE_rows.py - DO NOT EDIT BY HAND.") + + for cat_idx, (category, methods) in enumerate(CATEGORIES): + if cat_idx > 0: + lines.append(r"\midrule") + + category_tex = latex_escape(category) + lines.append(rf"\multicolumn{{6}}{{@{{}}l}}{{\emph{{{category_tex}}}}} \\") + + for method in methods: + display = latex_escape(DISPLAY_NAMES[method]) + values = aggregate_for_method(cell_values, method, metric) + formatted = [fmt_metric(metric, v) for v in values] + lines.append( + f"{display} & {formatted[0]} & {formatted[1]} & " + f"{formatted[2]} & {formatted[3]} & {formatted[4]} \\\\" + ) + + return lines + + +def write_fragments(cell_values: Dict[Tuple[str, str, str], float]) -> None: + os.makedirs(OUT_DIR, exist_ok=True) + + for metric, filename in OUTPUT_FILES.items(): + path = os.path.join(OUT_DIR, filename) + lines = render_metric_fragment(cell_values, metric) + with open(path, "w", encoding="utf-8", newline="\n") as f: + f.write("\n".join(lines)) + f.write("\n") + print(f"Wrote {path} ({len(lines)} lines)") + + +def sanity_check_methods() -> None: + flat_methods = [m for _, methods in CATEGORIES for m in methods] + if len(flat_methods) != 31: + raise ValueError(f"Expected 31 methods, got {len(flat_methods)}") + if len(set(flat_methods)) != 31: + dupes = sorted({m for m in flat_methods if flat_methods.count(m) > 1}) + raise ValueError(f"Duplicate methods in inventory: {dupes}") + missing_display = [m for m in flat_methods if m not in DISPLAY_NAMES] + if missing_display: + raise ValueError(f"Missing display names: {missing_display}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "csv", + nargs="?", + default=DEFAULT_INPUT_CSV, + help="Input long-format CSV: method_key,scene,metric,value", + ) + args = parser.parse_args() + + sanity_check_methods() + cell_values = load_long_csv(args.csv) + write_fragments(cell_values) + + print( + "Generated 3 fragments: 31 rows + 4 category dividers + 1 header " + "multicolumn each = 36 lines per file" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/generate_appendix_rows.py b/tools/generate_appendix_rows.py new file mode 100644 index 0000000000000000000000000000000000000000..566305746fbf1c6c011e6d0161471ab4910de829 --- /dev/null +++ b/tools/generate_appendix_rows.py @@ -0,0 +1,216 @@ +from __future__ import annotations + +import sys +import os +from pathlib import Path + +import pandas as pd + +# ----------------------------------------------------------------------------- +# 稳健路径配置 +# ----------------------------------------------------------------------------- +PROJECT_ROOT = Path("/root/autodl-tmp/SplatAtlas") + +# App C 候选输入 +APP_C_SRC_FULL = PROJECT_ROOT / "outputs" / "phase2" / "task_2_3_appE_full.csv" +APP_C_SRC_FILT = PROJECT_ROOT / "outputs" / "phase2" / "task_2_3_appE_filtered_P01_P04.csv" + +# App D 输入 +APP_D_SRC = PROJECT_ROOT / "outputs" / "phase4" / "task_4_1_seed_variance_summary_reconstructed.csv" + +OUT_DIR = PROJECT_ROOT / "tex" +OUT_DIR.mkdir(parents=True, exist_ok=True) +APP_C_OUT = OUT_DIR / "appC_paired_rows.tex" +APP_D_OUT = OUT_DIR / "appD_seed_summary_rows.tex" + +SEED_P95 = 0.384 + +# ----------------------------------------------------------------------------- +# 映射字典 +# ----------------------------------------------------------------------------- +VARIANT_DISPLAY = { + "pgsr": "PGSR", + "erankgs": "eRankGS", + "lightgaussian": "LightGaussian", + "steepgs": "SteepGS", +} + +METHOD_DISPLAY = { + "vanilla_3dgs": r"\texttt{vanilla\_3dgs}", + "analyticsplatting": r"\texttt{analyticsplatting}", + "erankgs": r"\texttt{erankgs}", + "ges": r"\texttt{ges}", + "lightgaussian": r"\texttt{lightgaussian}", + "minisplatting": r"\texttt{minisplatting}", + "opti3dgs": r"\texttt{opti3dgs}", + "pgsr": r"\texttt{pgsr}", + "steepgs": r"\texttt{steepgs}", + "3dgsmcmc": r"\texttt{3dgsmcmc}", +} + +SCENE_DISPLAY = { + "bicycle": "Bicycle", "bonsai": "Bonsai", "counter": "Counter", + "flowers": "Flowers", "garden": "Garden", "kitchen": "Kitchen", + "room": "Room", "stump": "Stump", "treehill": "Treehill", + "auditorium": "Auditorium", "ballroom": "Ballroom", "barn": "Barn", + "caterpillar": "Caterpillar", "courtroom": "Courtroom", "lighthouse": "Lighthouse", + "museum": "Museum", "palace": "Palace", "playground": "Playground", + "temple": "Temple", "train": "Train", "truck": "Truck", + "drjohnson": "DrJohnson", "playroom": "Playroom", + "chair": "Chair", "drums": "Drums", "ficus": "Ficus", "hotdog": "Hotdog", + "lego": "Lego", "materials": "Materials", "mic": "Mic", "ship": "Ship", +} + +def fmt_scene(key) -> str: + k = str(key).strip().lower() + name = SCENE_DISPLAY.get(k, str(key)) + return rf"\textsc{{{name}}}" + +def fmt_signed(x, prec: int = 3) -> str: + if pd.isna(x): return "---" + return f"{x:+.{prec}f}" + +def fmt_unsigned(x, prec: int = 3) -> str: + if pd.isna(x): return "---" + return f"{x:.{prec}f}" + +# ----------------------------------------------------------------------------- +# 生成逻辑 +# ----------------------------------------------------------------------------- +def generate_app_c() -> int: + if APP_C_SRC_FULL.exists(): + print(f"[App C] reading {APP_C_SRC_FULL}") + df = pd.read_csv(APP_C_SRC_FULL) + print(f"[App C] loaded {len(df)} rows total from FULL CSV") + df = df[df["pair_id"].isin(["P01", "P02", "P03", "P04"])].copy() + print(f"[App C] after P01-P04 filter: {len(df)} rows (expect 124)") + elif APP_C_SRC_FILT.exists(): + print(f"[App C] reading {APP_C_SRC_FILT}") + df = pd.read_csv(APP_C_SRC_FILT) + print(f"[App C] loaded {len(df)} rows total from FILTERED CSV (expect 124)") + else: + print("ERROR: missing App C source files. Searching for candidates:") + os.system('find /root/autodl-tmp/SplatAtlas/outputs -name "task_2_3_appE*.csv"') + sys.exit(1) + + df["_scene_disp"] = df["scene"].astype(str).str.lower().map(SCENE_DISPLAY).fillna(df["scene"]) + df = df.sort_values(["pair_id", "_scene_disp"]).reset_index(drop=True) + + lines: list[str] = [] + lines.append("% AUTOGENERATED by tools/generate_appendix_rows.py - DO NOT EDIT BY HAND.\n") + + last_pair = None + n_data_rows = 0 + for _, row in df.iterrows(): + if last_pair is not None and row["pair_id"] != last_pair: + lines.append(r"\midrule" + "\n") + last_pair = row["pair_id"] + + variant_disp = VARIANT_DISPLAY.get(row["variant"], row["variant"]) + scene_disp = fmt_scene(row["scene"]) + + cells = [ + row["pair_id"], + variant_disp, + scene_disp, + fmt_signed(row["delta_psnr_full"]), + fmt_signed(row["delta_sh_net_effect"]), + fmt_signed(row["delta_sh_corruption_rate"]), + fmt_signed(row["delta_opacity_net_effect"]), + fmt_signed(row["delta_opacity_pathology_rate"]), + fmt_signed(row["delta_coverage_error_fraction"]), + fmt_signed(row["delta_residual_error"]), + ] + lines.append(" & ".join(cells) + r" \\" + "\n") + n_data_rows += 1 + + APP_C_OUT.write_text("".join(lines)) + print(f"[App C] wrote {APP_C_OUT} ({n_data_rows} data rows)") + return n_data_rows + +def generate_app_d() -> int: + if not APP_D_SRC.exists(): + print("ERROR: missing App D source files. Searching for candidates:") + os.system('find /root/autodl-tmp/SplatAtlas/outputs -name "*seed_variance*summary*.csv"') + sys.exit(1) + + print(f"[App D] reading {APP_D_SRC}") + df = pd.read_csv(APP_D_SRC) + print(f"[App D] loaded {len(df)} rows total") + + scene_col = "SceneNormalized" if "SceneNormalized" in df.columns else "Scene" + method_col = "Method" + + out_lines: list[str] = [] + out_lines.append("% AUTOGENERATED by tools/generate_appendix_rows.py - DO NOT EDIT BY HAND.\n") + + n_data_rows = 0 + for scene_idx, scene in enumerate(["bonsai", "lego"]): + if scene_idx > 0: + out_lines.append(r"\midrule" + "\n") + + df_scene = df[df[scene_col].astype(str).str.lower() == scene].copy() + df_scene = df_scene.sort_values(method_col) + scene_disp = fmt_scene(scene) + + for _, row in df_scene.iterrows(): + method_key = row[method_col] + escaped_key = str(method_key).replace("_", r"\_") + method_disp = METHOD_DISPLAY.get(method_key, rf"\texttt{{{escaped_key}}}") + + psnr_mean = row["PSNR_mean"] + psnr_std = row["PSNR_std"] + psnr_min = row["PSNR_min"] + psnr_max = row["PSNR_max"] + status = str(row.get("Status", "OK")).strip().upper() + + if status == "METHOD_FAILURE": + flag = r"\textsc{Failure}" + elif status == "OUTLIER" or (not pd.isna(psnr_std) and psnr_std > SEED_P95): + flag = r"\textsc{Outlier}" + else: + flag = "---" + + cells = [ + method_disp, scene_disp, + fmt_unsigned(psnr_mean), fmt_unsigned(psnr_std), + fmt_unsigned(psnr_min), fmt_unsigned(psnr_max), + flag, + ] + out_lines.append(" & ".join(cells) + r" \\" + "\n") + n_data_rows += 1 + + APP_D_OUT.write_text("".join(out_lines)) + print(f"[App D] wrote {APP_D_OUT} ({n_data_rows} data rows)") + return n_data_rows + +# ----------------------------------------------------------------------------- +# Main & Sanity Checks +# ----------------------------------------------------------------------------- +if __name__ == "__main__": + print("\n========================================================") + print("1. Generating LaTeX row fragments") + print("========================================================") + c_rows = generate_app_c() + d_rows = generate_app_d() + + print("\n========================================================") + print("2. SANITY CHECKS") + print("========================================================") + print(f"[Check] App C output rows equals 124: {'PASS' if c_rows == 124 else 'FAIL'} ({c_rows})") + print(f"[Check] App D output rows equals 20: {'PASS' if d_rows == 20 else 'FAIL'} ({d_rows})") + + # 验证 App D 中的特殊标记 + with open(APP_D_OUT, "r") as f: + d_content = f.read() + ana_outlier = ("analyticsplatting" in d_content and "Outlier" in d_content) + era_failure = ("erankgs" in d_content and "Failure" in d_content) + print(f"[Check] analyticsplatting has \\textsc{{Outlier}}: {'PASS' if ana_outlier else 'FAIL'}") + print(f"[Check] erankgs has \\textsc{{Failure}}: {'PASS' if era_failure else 'FAIL'}") + + print("\n========================================================") + print("Generated:") + print(f"- {APP_C_OUT}: {c_rows} rows") + print(f"- {APP_D_OUT}: {d_rows} rows") + print("========================================================") + diff --git a/ufd_evalkit/__init__.py b/ufd_evalkit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/ufd_evalkit/alignment.py b/ufd_evalkit/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..230512a63fcb1c88378b03b0e2c7550aa75a1aad --- /dev/null +++ b/ufd_evalkit/alignment.py @@ -0,0 +1,101 @@ +import os +import json +import numpy as np +from scipy.spatial import cKDTree +from plyfile import PlyData # 仅用于读取 COLMAP 的标准真实点云 + +def umeyama_transform(P, Q): + """标准的 Umeyama 算法求解 7-DoF (s, R, t) 变换矩阵""" + mu_P = P.mean(axis=0) + mu_Q = Q.mean(axis=0) + P_centered = P - mu_P + Q_centered = Q - mu_Q + + C = P_centered.T @ Q_centered / P.shape[0] + U, S, Vt = np.linalg.svd(C) + + R = Vt.T @ U.T + if np.linalg.det(R) < 0: + Vt[2, :] *= -1 + R = Vt.T @ U.T + + var_P = np.var(P, axis=0).sum() + if var_P < 1e-8: + scale = 1.0 + else: + scale = (1.0 / var_P) * np.sum(S) + + t = mu_Q - scale * (R @ mu_P) + + T = np.eye(4) + T[:3, :3] = scale * R + T[:3, 3] = t + return T + +def icp_umeyama_align(source_pts, target_pts, max_iters=20, tolerance=1e-5): + """使用 ICP 匹配点对，使用 Umeyama 求解矩阵""" + if len(source_pts) > 50000: + idx = np.random.choice(len(source_pts), 50000, replace=False) + src = source_pts[idx].copy() + else: + src = source_pts.copy() + + target_tree = cKDTree(target_pts) + T_accum = np.eye(4) + prev_dist = float('inf') + + for i in range(max_iters): + dists, indices = target_tree.query(src) + matched_target = target_pts[indices] + + T = umeyama_transform(src, matched_target) + src_homo = np.hstack((src, np.ones((src.shape[0], 1)))) + src = (T @ src_homo.T).T[:, :3] + T_accum = T @ T_accum + + if i > 0 and np.mean(dists) - prev_dist > -tolerance: + break + prev_dist = np.mean(dists) + + return T_accum + +def assess_manifold_collapse(model, model_dir, colmap_ply_path=None): + try: + # 优雅降级：如果模型没有提供离散空间点（比如隐式表征），直接跳过此项检测 + if not hasattr(model, 'get_spatial_centers'): + return {"manifold_collapse": False, "reason": "Model lacks explicit discrete centers"} + + try: + xyz_pred = model.get_spatial_centers().cpu().numpy() + except NotImplementedError: + return {"manifold_collapse": False, "reason": "Method not implemented"} + + if len(xyz_pred) < 100: + return {"manifold_collapse": True, "reason": "Too few points"} + + # --- 拓扑学流形坍塌检测 --- + cov = np.cov(xyz_pred.T) + eigenvalues, _ = np.linalg.eigh(cov) + if eigenvalues[0] < 1e-6: + return {"manifold_collapse": True, "reason": "Zero variance in one dimension (Flat collapse)"} + + # --- 锚点绝对坐标系对齐 (ICP-Umeyama) --- + matrix = np.eye(4) + if colmap_ply_path and os.path.exists(colmap_ply_path): + try: + colmap_data = PlyData.read(colmap_ply_path) + cv = colmap_data.elements[0].data + xyz_colmap = np.vstack([cv['x'], cv['y'], cv['z']]).T + + print(f" [Alignment] 正在对齐到绝对宇宙 (ICP)... 锚点数量: {len(xyz_colmap)}") + matrix = icp_umeyama_align(xyz_pred, xyz_colmap) + except Exception as e: + print(f" [Alignment] 锚点对齐失败 ({e})，回退为单位矩阵。") + + matrix_path = os.path.join(model_dir, "alignment_matrix.json") + with open(matrix_path, 'w') as f: + json.dump(matrix.tolist(), f, indent=4) + + return {"manifold_collapse": False, "reason": "Normal"} + except Exception as e: + return {"manifold_collapse": False, "reason": f"Detection skipped due to error: {str(e)}"} diff --git a/ufd_evalkit/geometric.py b/ufd_evalkit/geometric.py new file mode 100644 index 0000000000000000000000000000000000000000..f2c31d24bd6c8b6ff7d9dea198c731c50b847486 --- /dev/null +++ b/ufd_evalkit/geometric.py @@ -0,0 +1,52 @@ +import os +import json +import numpy as np +import torch + +def compute_geometric_pathology(model, cameras=None) -> dict: + """ + 提取几何病理学指纹。 + 底层计算已通过 IoC 模式下放给各个 Wrapper 的 compute_physical_metrics 接口。 + 此处仅负责集中管控诊断阈值 (Pathology Flags)。 + """ + print(f"🔬 [UFD-EvalKit] 正在请求变体自述几何资产...") + + try: + metrics = model.compute_physical_metrics(cameras=cameras) + except NotImplementedError: + print(" ⚠️ [UFD-EvalKit] Wrapper 未实现 compute_physical_metrics，返回空字典。") + metrics = {} + except Exception as e: + return {"error": f"提取物理指标失败: {str(e)}"} + + # 集中研判诊断标签 (Pathology Flags)，保持全榜单阈值标准绝对统一 + flags = [] + if metrics.get("gamma_median", 0) > 18.0: + flags.append("SEVERE_ANISOTROPY_DISTORTION") + if metrics.get("alpha_mean", 1.0) < 0.05: + flags.append("OPACITY_MINIMIZATION_COLLAPSE") + if metrics.get("scale_mean", 0) > 5.0: + flags.append("GLOBAL_SCALE_INFLATION") + + metrics["pathology_flags"] = flags + return metrics + +def depth_to_normal(depth): + """从深度图反推屏幕空间法向 (Method-agnostic)""" + if len(depth.shape) == 2: depth = depth.unsqueeze(0).unsqueeze(0) + elif len(depth.shape) == 3: depth = depth.unsqueeze(0) + + # 使用简单差分计算深度梯度 + dx = depth[:, :, :, 1:] - depth[:, :, :, :-1] + dy = depth[:, :, 1:, :] - depth[:, :, :-1, :] + + # 填充边界对齐尺寸 + dx = torch.nn.functional.pad(dx, (0, 1, 0, 0)) + dy = torch.nn.functional.pad(dy, (0, 0, 0, 1)) + + # 构造近似法向量并归一化 + normal = torch.cat([-dx, -dy, torch.ones_like(depth)], dim=1) + normal = torch.nn.functional.normalize(normal, dim=1) + + # 映射到 [0, 1] 用于可视化存储 + return (normal + 1.0) / 2.0 diff --git a/ufd_evalkit/photometric.py b/ufd_evalkit/photometric.py new file mode 100644 index 0000000000000000000000000000000000000000..37fdb594de1d7ec32f4a7d81774499e24b3f0915 --- /dev/null +++ b/ufd_evalkit/photometric.py @@ -0,0 +1,80 @@ +import os +import torch +import numpy as np +from PIL import Image +import lpips +from skimage.metrics import structural_similarity as compute_ssim + + +def _load_rgb(path: str) -> np.ndarray: + """Load any PNG as RGB uint8. Composites RGBA over white if needed.""" + img = Image.open(path) + if img.mode == 'RGBA': + bg = Image.new('RGB', img.size, (255, 255, 255)) + bg.paste(img, mask=img.split()[3]) + img = bg + else: + img = img.convert('RGB') + return np.array(img) + + +def compute_photometric_metrics(render_dir: str, gt_dir: str) -> dict: + loss_fn_vgg = lpips.LPIPS(net='vgg').to("cuda") + psnr_list, lpips_list, ssim_list = [], [], [] + skipped = 0 + + render_files = sorted(f for f in os.listdir(render_dir) if f.endswith('.png')) + for img_name in render_files: + render_path = os.path.join(render_dir, img_name) + gt_path = os.path.join(gt_dir, img_name) + if not os.path.exists(gt_path): + skipped += 1 + continue + + try: + r_np = _load_rgb(render_path) + g_np = _load_rgb(gt_path) + if r_np.shape != g_np.shape: + g_pil = Image.fromarray(g_np).resize( + (r_np.shape[1], r_np.shape[0]), Image.LANCZOS) + g_np = np.array(g_pil) + + render_img = torch.from_numpy(r_np).float().cuda() / 255.0 + gt_img = torch.from_numpy(g_np).float().cuda() / 255.0 + render_img = render_img.permute(2, 0, 1).unsqueeze(0) + gt_img = gt_img.permute(2, 0, 1).unsqueeze(0) + + r8 = (torch.clamp(render_img, 0, 1) * 255).byte().float() / 255.0 + g8 = (torch.clamp(gt_img, 0, 1) * 255).byte().float() / 255.0 + + mse = torch.mean((r8 - g8) ** 2) + psnr = (-10.0 * torch.log10(mse)).item() if mse.item() > 1e-12 else 100.0 + with torch.no_grad(): + lpips_val = loss_fn_vgg(r8 * 2 - 1, g8 * 2 - 1).item() + + r_chw = r8.squeeze().permute(1, 2, 0).cpu().numpy() + g_chw = g8.squeeze().permute(1, 2, 0).cpu().numpy() + ssim_val = compute_ssim(r_chw, g_chw, data_range=1.0, channel_axis=-1) + + if not (np.isfinite(psnr) and np.isfinite(lpips_val) and np.isfinite(ssim_val)): + print(f"[WARN] non-finite on {img_name}: PSNR={psnr} LPIPS={lpips_val} SSIM={ssim_val}") + skipped += 1 + continue + + psnr_list.append(psnr) + lpips_list.append(lpips_val) + ssim_list.append(ssim_val) + except Exception as e: + print(f"[ERROR] {img_name}: {e}") + skipped += 1 + + n = len(psnr_list) + print(f"[photometric] {render_dir}: used {n} pairs, skipped {skipped}") + if n == 0: + return {"PSNR": float('nan'), "LPIPS": float('nan'), "SSIM": float('nan'), "n": 0} + return { + "PSNR": float(np.mean(psnr_list)), + "LPIPS": float(np.mean(lpips_list)), + "SSIM": float(np.mean(ssim_list)), + "n": n, + } diff --git a/ufd_evalkit/run_eval.py b/ufd_evalkit/run_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..d1beb968cf8912d5fca3306e7eaf7e29ac521224 --- /dev/null +++ b/ufd_evalkit/run_eval.py @@ -0,0 +1,114 @@ +import os +import json + +class NpEncoder(json.JSONEncoder): + def default(self, obj): + if hasattr(obj, 'item'): + return obj.item() + return super().default(obj) + +import argparse +import shutil +import numpy as np +import torch +import sys + +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) +from core.registry import METHOD_REGISTRY +import methods +from photometric import compute_photometric_metrics +from geometric import compute_geometric_pathology +from alignment import assess_manifold_collapse + +def load_cameras_from_json(cam_path): + if not os.path.exists(cam_path): return [] + with open(cam_path, 'r') as f: cams_data = json.load(f) + class DummyCam: + def __init__(self, center): + self.camera_center = torch.tensor(center, dtype=torch.float32) + return [DummyCam(c['center']) for c in cams_data if 'center' in c] + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--method", type=str, required=True) + parser.add_argument("--scene", type=str, required=True) + parser.add_argument("--render_dir", type=str, required=True) + parser.add_argument("--gt_dir", type=str, required=True) + parser.add_argument("--ply_path", type=str, required=True) # 保持参数兼容外层流水线 + parser.add_argument("--output_json", type=str, required=True) + parser.add_argument("--prior_type", type=str, default="none") + parser.add_argument("--depth_dir", type=str, default="") + parser.add_argument("--skip_quarantine", action="store_true") + parser.add_argument("--colmap_dir", type=str, default="") + return parser.parse_args() + +def main(): + args = parse_args() + os.makedirs(os.path.dirname(args.output_json), exist_ok=True) + + # 从 ply_path 中反向推导出 model_path 和 iteration，实现无缝过渡 + model_path = args.ply_path.split('/point_cloud/')[0] + try: + iteration = int(args.ply_path.split('/iteration_')[1].split('/')[0]) + except IndexError: + iteration = 30000 + + # 1. 实例化变体并加载资产 + model_class = methods.load_method(args.method) + _OUTDOOR_360 = {"bicycle", "flowers", "garden", "stump", "treehill"} + _res = 4 if args.scene in _OUTDOOR_360 else 2 + dataset_config = {"source_path": args.colmap_dir, "model_path": model_path, "resolution": _res} + model = model_class(dataset_config, hyperparams={}) + + try: + model.load(model_path, iteration) + except Exception as e: + print(f"❌ [Eval] 模型资产加载失败: {e}") + return + + cams = load_cameras_from_json(os.path.join(model_path, "cameras.json")) + + # 2. 调用重构后的几何病理学检测 + geo_metrics = compute_geometric_pathology(model, cameras=cams) + flags = geo_metrics.get("pathology_flags", []) + + # 3. 调用重构后的流形坍塌检测 + colmap_ply_path = os.path.join(args.colmap_dir, "sparse", "0", "points3D.ply") if args.colmap_dir else None + alignment_res = assess_manifold_collapse(model, model_path, colmap_ply_path=colmap_ply_path) + + if alignment_res.get("manifold_collapse", False): + flags.append("MANIFOLD_COLLAPSE") + pass + else: + del model; torch.cuda.empty_cache() + photo_metrics = compute_photometric_metrics(args.render_dir, args.gt_dir) + + if args.depth_dir and os.path.exists(args.depth_dir): + depth_vars = [np.var(np.load(os.path.join(args.depth_dir, f))[np.load(os.path.join(args.depth_dir, f)) > 0.05]) + for f in os.listdir(args.depth_dir) if f.endswith('.npy')] + mean_depth_var = np.mean(depth_vars) if depth_vars else 1.0 + geo_metrics["mean_depth_variance"] = round(float(mean_depth_var), 4) + if mean_depth_var < 0.01: + flags.append("FLAT_WALL_COLLAPSE") + + geo_metrics["pathology_flags"] = flags + + report = { + "method": args.method, + "scene": args.scene, + "prior_type": args.prior_type, + "photometric": photo_metrics, + "geometric": geo_metrics, + "pathology_flags": list(set(flags)) + } + + with open(args.output_json, 'w') as f: + json.dump(report, f, indent=4) + + if len(flags) > 0 and not args.skip_quarantine: + master_log_path = os.path.join(os.path.dirname(model_path), "pathology_master_log.txt") + with open(master_log_path, "a") as log_f: + log_f.write(f"Method: {args.method} | Scene: {args.scene} | Path: {model_path} | Flags: {','.join(list(set(flags)))}\n") + +if __name__ == "__main__": + main()