File size: 5,257 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | #!/usr/bin/env python3
"""
Randomly pick N revisit candidates (gen-only MP4) and compute average PSNR(+LPIPS if available).
Intended for eval_v2/static_consistency outputs:
- <evals_root>/in_domain/loop_closure/**/*_gen_only.mp4
- <evals_root>/in_domain/combo_revisit_fixed_first/**/*_gen_only.mp4
Outputs:
- per_video metrics JSON
- aggregate summary JSON
- optional side-by-side visualization PNGs (first vs last frame)
"""
from __future__ import annotations
import argparse
import json
import os
import random
import sys
from typing import Any, Dict, List, Optional, Tuple
def _discover_candidates(evals_root: str) -> List[str]:
cands: List[str] = []
for root, _dirs, files in os.walk(evals_root):
for f in files:
if f.endswith("_gen_only.mp4"):
cands.append(os.path.join(root, f))
return sorted(cands)
def _safe_mean(values: List[float]) -> Optional[float]:
if not values:
return None
return float(sum(values) / len(values))
def main() -> None:
p = argparse.ArgumentParser(description="Aggregate revisit metrics by random sampling")
p.add_argument("--evals_root", type=str, required=True, help="Path to static_consistency/in_domain (or similar)")
p.add_argument("--num_samples", type=int, default=5, help="How many videos to sample (min(len, num_samples))")
p.add_argument("--seed", type=int, default=42)
p.add_argument("--device", type=str, default="cuda")
p.add_argument("--output_dir", type=str, required=True)
p.add_argument("--write_viz", action="store_true", help="Also export side-by-side PNG for sampled videos")
p.add_argument("--viz_dir", type=str, default=None, help="Override visualization output directory")
args = p.parse_args()
evals_root = os.path.abspath(args.evals_root)
output_dir = os.path.abspath(args.output_dir)
os.makedirs(output_dir, exist_ok=True)
candidates = _discover_candidates(evals_root)
if not candidates:
summary = {
"num_candidates": 0,
"num_samples": 0,
"aggregate": {},
"per_video": [],
"note": "No *_gen_only.mp4 found under evals_root",
}
with open(os.path.join(output_dir, "revisit_metrics_summary.json"), "w", encoding="utf-8") as f:
json.dump(summary, f, indent=2)
print(f"[aggregate_revisit_metrics] No candidates under: {evals_root}", file=sys.stderr)
return
rng = random.Random(args.seed)
num_samples = min(len(candidates), max(1, int(args.num_samples)))
selected = rng.sample(candidates, k=num_samples) if num_samples < len(candidates) else candidates
# Import metrics with local sys.path insertion to avoid needing package __init__.py
_script_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _script_dir)
import psnr_lpips # type: ignore
# Optional visualization
if args.write_viz:
viz_dir = args.viz_dir or os.path.join(output_dir, "viz_first_vs_last")
os.makedirs(viz_dir, exist_ok=True)
else:
viz_dir = None
per_video: List[Dict[str, Any]] = []
psnr_list: List[float] = []
lpips_list: List[float] = []
for vp in selected:
rel = os.path.relpath(vp, evals_root)
try:
res = psnr_lpips.compute_revisit_metrics(vp, device=args.device)
per_video.append({"rel": rel, **res})
if res.get("psnr") is not None:
psnr_list.append(float(res["psnr"]))
if res.get("lpips") is not None:
lpips_list.append(float(res["lpips"]))
except Exception as e:
per_video.append({"rel": rel, "error": str(e)})
if args.write_viz and viz_dir is not None:
# Export side-by-side png
try:
import subprocess
_viz_script = os.path.join(os.path.dirname(_script_dir), "visualize", "revisit_pairs_viz.py")
out_png = os.path.join(viz_dir, rel.replace(os.sep, "__") + ".png")
os.makedirs(os.path.dirname(out_png), exist_ok=True)
subprocess.run(
["python3", _viz_script, "--video", vp, "--output", out_png],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
except Exception:
pass
aggregate = {
"num_candidates": len(candidates),
"num_samples": len(selected),
"mean_psnr": _safe_mean(psnr_list),
"mean_lpips": _safe_mean(lpips_list),
"note": "LPIPS is None for each video when lpips dependency is missing.",
}
summary = {
"evals_root": evals_root,
"seed": args.seed,
"aggregate": aggregate,
"per_video": per_video,
}
with open(os.path.join(output_dir, "revisit_metrics_summary.json"), "w", encoding="utf-8") as f:
json.dump(summary, f, indent=2)
with open(os.path.join(output_dir, "revisit_metrics_per_video.json"), "w", encoding="utf-8") as f:
json.dump(per_video, f, indent=2)
print(f"[aggregate_revisit_metrics] wrote: {output_dir}/revisit_metrics_summary.json")
if __name__ == "__main__":
main()
|