| |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import imageio.v2 as imageio |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Validate an indexed Self-Forcing video evaluation set." |
| ) |
| parser.add_argument("video_dir", type=Path) |
| parser.add_argument("--expected-videos", type=int, default=100) |
| parser.add_argument("--expected-frames", type=int, default=81) |
| parser.add_argument("--expected-height", type=int, default=480) |
| parser.add_argument("--expected-width", type=int, default=832) |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if not args.video_dir.is_dir(): |
| raise FileNotFoundError(f"Video directory does not exist: {args.video_dir}") |
|
|
| expected_names = { |
| f"{index}-0_ema.mp4" for index in range(args.expected_videos) |
| } |
| video_paths = sorted(args.video_dir.glob("*.mp4")) |
| actual_names = {path.name for path in video_paths} |
| if actual_names != expected_names: |
| missing = sorted(expected_names - actual_names) |
| extra = sorted(actual_names - expected_names) |
| raise ValueError( |
| f"Video set mismatch: count={len(video_paths)}, " |
| f"missing={missing[:10]}, extra={extra[:10]}" |
| ) |
|
|
| for path in video_paths: |
| reader = imageio.get_reader(path, format="ffmpeg") |
| try: |
| frame_count = reader.count_frames() |
| if frame_count != args.expected_frames: |
| raise ValueError( |
| f"{path.name}: expected {args.expected_frames} frames, " |
| f"found {frame_count}" |
| ) |
| first_frame = reader.get_data(0) |
| last_frame = reader.get_data(args.expected_frames - 1) |
| finally: |
| reader.close() |
| expected_shape = (args.expected_height, args.expected_width, 3) |
| if first_frame.shape != expected_shape or last_frame.shape != expected_shape: |
| raise ValueError( |
| f"{path.name}: expected frame shape {expected_shape}, " |
| f"found first={first_frame.shape}, last={last_frame.shape}" |
| ) |
|
|
| print( |
| f"valid video set: path={args.video_dir.resolve()} " |
| f"videos={len(video_paths)} frames_per_video={args.expected_frames} " |
| f"resolution={args.expected_width}x{args.expected_height}" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
