File size: 2,433 Bytes
2847d0b | 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 | #!/usr/bin/env python3
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()
|