#!/usr/bin/env python3 """Multi-scene benchmark driver. Each scene runs as its own `spirula train` process, so the engine (world splats, bilagrid/PPISP/background, optimizer moments, colour-space matrices, device pool) starts clean. Several scenes in one process leak state between them and silently degrade the later ones. Each run writes its own metrics.json -- l1/psnr/ssim, their colour-corrected `cc_` variants, `training_time` and `engine_vram`, all native. This script reads those back and prints a table. With --lpips it also runs eval_lpips.py per scene, which needs torch + torchmetrics; without it, this script needs nothing beyond the standard library. python3 reference/python/benchmark.py 360_v2 --data /path/to/360_v2 python3 reference/python/benchmark.py zipnerf --data /path/to/zipnerf \\ --preset academic-baseline --lpips Extra arguments after `--` are appended to every `spirula train` command: python3 reference/python/benchmark.py 360_v2 --data ... -- --cap-max 3000000 """ from __future__ import annotations import argparse import datetime import json import shutil import subprocess import sys from pathlib import Path # (scene, downscale factor). The rounding mode differs per scene because the # published images_N folders were produced with different rounding. BENCHMARKS: dict[str, dict] = { "360_v2": { "help": "7 permissively released scenes from Mip-NeRF 360", "url": "http://storage.googleapis.com/gresearch/refraw360/360_v2.zip", "scenes": [("bicycle", 4), ("garden", 4), ("stump", 4), ("bonsai", 2), ("counter", 2), ("kitchen", 2), ("room", 2)], # garden's images_4 was rounded, the rest ceiled. "rounding": lambda scene: "round" if scene == "garden" else "ceil", "extra": lambda scene: [], }, "zipnerf": { "help": "4 fisheye scenes from Zip-NeRF", "url": "https://smerf-3d.github.io/#data", "scenes": [("alameda", 4), ("berlin", 4), ("london", 4), ("nyc", 4)], "rounding": lambda scene: "ceil", "extra": lambda scene: [], }, } def find_trainer() -> str: exe = shutil.which("spirula") if exe: return exe local = Path(__file__).resolve().parents[2] / "build" / "spirula" if local.is_file(): return str(local) raise SystemExit("`spirula` not found on PATH or in ./build -- build it first") def run_scene(trainer: str, preset: str, data_dir: Path, out_prefix: Path, out_name: str, downscale: int, rounding: str, extra: list[str], want_lpips: bool) -> dict | None: metrics_path = out_prefix / out_name / "metrics.json" # A stale file from a previous run with the same name would masquerade as # this run's result. metrics_path.unlink(missing_ok=True) cmd = [ trainer, "train", preset, "--data", str(data_dir), "--data-format", "colmap", "--colmap-recon-dir", "sparse/0", "--rescale-camera-to-fit", str(downscale), "--image-dir", f"images_{downscale}", "--downscale-rounding-mode", rounding, "--eval-mode", "interval", "--save-eval-images", "1", "--load-depths", "0", "--load-normals", "0", "--steps-per-save", "0", "--disable-viewer", "1", "--keep-viewer-alive", "0", "--output-dir-prefix", str(out_prefix), "--output-dir-name", out_name, ] + extra print(">>>", " ".join(cmd), flush=True) if subprocess.call(cmd) != 0: print(f"!!! {out_name} exited non-zero", flush=True) return None if not metrics_path.exists(): print(f"!!! {out_name} wrote no metrics.json", flush=True) return None if want_lpips: lpips = Path(__file__).with_name("eval_lpips.py") rc = subprocess.call([sys.executable, str(lpips), str(out_prefix / out_name), "--write"]) if rc != 0: print(f"!!! LPIPS failed for {out_name}; metrics are still valid " f"without it", flush=True) return json.loads(metrics_path.read_text()) def print_table(results: list[tuple[str, dict | None]]) -> None: ok = [(name, m) for name, m in results if m is not None] failed = [name for name, m in results if m is None] if not ok: print("\nEvery scene failed.") return # Scalar keys only; avg_* is what a scene summarizes to. keys: list[str] = [] for _, m in ok: for k, v in m.items(): if isinstance(v, list) or k in keys: continue keys.append(k) w = max(len(k) for k in keys) + 2 print("\n" + "metric".ljust(w) + "".join(n[:14].rjust(16) for n, _ in ok)) for k in keys: row = k[len("avg_"):] if k.startswith("avg_") else k line = row.ljust(w) for _, m in ok: v = m.get(k) line += ("" if v is None else f"{v:.5g}" if isinstance(v, float) else str(v)).rjust(16) print(line) if failed: print("\nfailed:", ", ".join(failed)) def main() -> None: ap = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("benchmark", choices=sorted(BENCHMARKS)) ap.add_argument("--data", type=Path, required=True, help="folder holding the benchmark's scenes") ap.add_argument("--preset", default="academic-baseline", help="spirula train preset (default: academic-baseline)") ap.add_argument("--output-dir-prefix", type=Path, default=Path("outputs")) ap.add_argument("--lpips", action="store_true", help="also run eval_lpips.py per scene (needs torchmetrics)") ap.add_argument("rest", nargs="*", help="after `--`, extra flags for every spirula train call") args = ap.parse_args() spec = BENCHMARKS[args.benchmark] if not args.data.is_dir(): raise SystemExit(f"dataset not found at {args.data}\n see {spec['url']}") trainer = find_trainer() tag = (f"{args.benchmark}-{args.preset}-" f"{datetime.datetime.now().strftime('%Y%m%d-%H%M%S')}") results = [] for scene, downscale in spec["scenes"]: print(f"\n=== {scene} ===", flush=True) results.append((scene, run_scene( trainer, args.preset, args.data / scene, args.output_dir_prefix, f"benchmark-{tag}-{scene}", downscale, spec["rounding"](scene), spec["extra"](scene) + list(args.rest), args.lpips))) print_table(results) if __name__ == "__main__": main()