93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
"""Score a folder of videos in parallel across multiple GPUs.
|
|
|
|
Uses :meth:`Evaluator.evaluate(samples=[...])`, which round-robins each
|
|
sample dict across the GPU replicas the evaluator was built with.
|
|
|
|
Example::
|
|
|
|
python examples/inference/eval/score_folder.py \\
|
|
--videos generated/ \\
|
|
--metrics vbench.aesthetic_quality,vbench.subject_consistency \\
|
|
--num-gpus 4 \\
|
|
--output scores.json
|
|
|
|
Pair each generated video with a same-name reference video (e.g.
|
|
``ref/<stem>.mp4``) by passing ``--reference-dir``::
|
|
|
|
python examples/inference/eval/score_folder.py \\
|
|
--videos generated/ --reference-dir ref/ \\
|
|
--metrics common.psnr,common.ssim,common.lpips \\
|
|
--num-gpus 4
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
from fastvideo.eval import create_evaluator
|
|
|
|
|
|
def _list_videos(directory: Path) -> list[Path]:
|
|
exts = {".mp4", ".avi", ".mov", ".mkv", ".gif"}
|
|
return sorted(p for p in directory.iterdir() if p.suffix.lower() in exts)
|
|
|
|
|
|
def main() -> None:
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("--videos", type=Path, required=True, help="Directory of generated videos.")
|
|
p.add_argument("--reference-dir",
|
|
type=Path,
|
|
default=None,
|
|
help="Directory of reference videos with matching stems.")
|
|
p.add_argument("--metrics", default="vbench.aesthetic_quality")
|
|
p.add_argument("--num-gpus", type=int, default=1)
|
|
p.add_argument("--fps", type=float, default=None, help="Frame-rate annotation for fps-aware metrics.")
|
|
p.add_argument("--output", type=Path, default=Path("scores.json"))
|
|
args = p.parse_args()
|
|
|
|
video_paths = _list_videos(args.videos)
|
|
if not video_paths:
|
|
raise SystemExit(f"No videos under {args.videos}")
|
|
print(f"Found {len(video_paths)} videos in {args.videos}")
|
|
|
|
metrics: list[str] | str = (args.metrics
|
|
if args.metrics == "all" else [m.strip() for m in args.metrics.split(",") if m.strip()])
|
|
evaluator = create_evaluator(metrics=metrics, num_gpus=args.num_gpus)
|
|
|
|
# Build per-video sample dicts holding *paths*, not pre-loaded
|
|
# tensors. Each path is decoded inside the worker thread that picks
|
|
# up its sample, so peak resident memory is bounded by num_gpus
|
|
# rather than scaling with the size of the folder.
|
|
samples: list[dict] = []
|
|
for vp in video_paths:
|
|
sample: dict = {"video": str(vp)}
|
|
if args.reference_dir is not None:
|
|
ref_path = args.reference_dir / vp.name
|
|
if not ref_path.is_file():
|
|
raise FileNotFoundError(f"Missing reference for {vp.name} at {ref_path}")
|
|
sample["reference"] = str(ref_path)
|
|
if args.fps is not None:
|
|
sample["fps"] = args.fps
|
|
samples.append(sample)
|
|
|
|
print(f"Scoring with {len(evaluator.metric_names)} metric(s) "
|
|
f"on {evaluator.num_gpus} GPU(s)...")
|
|
all_results = evaluator.evaluate(samples=samples)
|
|
evaluator.shutdown()
|
|
|
|
payload = [{
|
|
"video": str(vp),
|
|
"scores": {
|
|
name: r.score
|
|
for name, r in results.items()
|
|
},
|
|
} for vp, results in zip(video_paths, all_results)]
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(payload, indent=2))
|
|
print(f"Wrote {args.output}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|