scripts/evaluate_long_video_vbench.py

#!/usr/bin/env python3
"""Run all standard VBench dimensions for one long-video condition."""

from __future__ import annotations

import argparse
import os
import tempfile
from pathlib import Path


def preparse_gpu() -> str:
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument("--gpu", required=True)
    args, _ = parser.parse_known_args()
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    os.environ.setdefault("MPLCONFIGDIR", tempfile.mkdtemp(prefix="vbench_mpl_"))
    return args.gpu


GPU = preparse_gpu()

from vbench import VBench


DIMENSIONS = [
    "subject_consistency", "background_consistency", "motion_smoothness",
    "aesthetic_quality", "imaging_quality",
]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gpu", default=GPU)
    parser.add_argument("--video_dir", type=Path, required=True)
    parser.add_argument("--output_dir", type=Path, required=True)
    parser.add_argument("--name", required=True)
    args = parser.parse_args()
    args.video_dir = args.video_dir.resolve()
    args.output_dir = args.output_dir.resolve()
    args.output_dir.mkdir(parents=True, exist_ok=True)
    bench = VBench(
        device="cuda", full_info_dir=str(args.video_dir / "full_info.json"),
        output_path=str(args.output_dir),
    )
    bench.evaluate(
        videos_path=str(args.video_dir), name=args.name,
        dimension_list=DIMENSIONS, mode="custom_input",
    )


if __name__ == "__main__":
    main()