File size: 3,859 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 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 | #!/usr/bin/env python3
"""Build VBench custom-input directories from dynamic-gate video outputs."""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
DIMENSIONS = [
"subject_consistency",
"background_consistency",
"motion_smoothness",
"aesthetic_quality",
"imaging_quality",
]
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--dataset_root", type=Path, required=True)
parser.add_argument("--output_root", type=Path, required=True)
parser.add_argument(
"--condition",
action="append",
nargs=2,
metavar=("NAME", "VIDEO_DIR"),
required=True,
help="Condition name and directory containing prompt_XXXX.mp4 files.",
)
parser.add_argument("--prompt_start", type=int, default=90)
parser.add_argument("--prompt_end", type=int, default=100)
parser.add_argument(
"--combined_name",
default=None,
help="Also build one flattened input directory for a single shared VBench run.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
dataset_root = args.dataset_root.resolve()
output_root = args.output_root.resolve()
output_root.mkdir(parents=True, exist_ok=True)
combined_dir = output_root / args.combined_name if args.combined_name else None
combined_info = []
if combined_dir is not None:
combined_dir.mkdir(parents=True, exist_ok=True)
for name, source_text in args.condition:
source_dir = Path(source_text).resolve()
destination_dir = output_root / name
destination_dir.mkdir(parents=True, exist_ok=True)
full_info = []
for prompt_id in range(args.prompt_start, args.prompt_end):
filename = f"prompt_{prompt_id:04d}.mp4"
source = source_dir / filename
if not source.is_file():
raise FileNotFoundError(source)
destination = destination_dir / filename
if destination.is_symlink() or destination.exists():
destination.unlink()
destination.symlink_to(os.path.relpath(source, destination_dir))
metadata_path = dataset_root / f"prompt_{prompt_id:04d}" / "metadata.json"
metadata = json.loads(metadata_path.read_text(encoding="utf-8"))
full_info.append(
{
"video_list": [filename],
"prompt_en": metadata["prompt"],
"dimension": DIMENSIONS,
}
)
if combined_dir is not None:
combined_filename = f"{name}__{filename}"
combined_destination = combined_dir / combined_filename
if combined_destination.is_symlink() or combined_destination.exists():
combined_destination.unlink()
combined_destination.symlink_to(
os.path.relpath(source, combined_dir)
)
combined_info.append(
{
"video_list": [combined_filename],
"prompt_en": metadata["prompt"],
"dimension": DIMENSIONS,
}
)
(destination_dir / "full_info.json").write_text(
json.dumps(full_info, indent=2) + "\n", encoding="utf-8"
)
print(f"[prepared] {name}: {len(full_info)} videos -> {destination_dir}")
if combined_dir is not None:
(combined_dir / "full_info.json").write_text(
json.dumps(combined_info, indent=2) + "\n", encoding="utf-8"
)
print(
f"[prepared] combined: {len(combined_info)} videos -> {combined_dir}"
)
if __name__ == "__main__":
main()
|