| |
| """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() |
|
|
