File size: 1,517 Bytes
bc29ee3 | 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 | #!/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()
|