File size: 1,511 Bytes
ebec118 | 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 | """Decode training videos to 15 fps frames (fit in 448 px) as uint8 arrays,
and mark scene cuts so no training triplet spans one."""
import json, subprocess, sys
from pathlib import Path
import numpy as np
src, out = Path(sys.argv[1]), Path(sys.argv[2])
out.mkdir(parents=True, exist_ok=True)
index = []
for video in sorted(src.glob("v*.webm")):
probe = json.loads(subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
"stream=width,height", "-of", "json", str(video)], capture_output=True, text=True).stdout)
w, h = probe["streams"][0]["width"], probe["streams"][0]["height"]
s = 448 / max(w, h)
W, H = int(round(w * s / 8) * 8), int(round(h * s / 8) * 8)
raw = subprocess.run(["ffmpeg", "-v", "error", "-i", str(video), "-vf", f"fps=15,scale={W}:{H}:flags=area",
"-f", "rawvideo", "-pix_fmt", "rgb24", "-"], capture_output=True).stdout
frames = np.frombuffer(raw, np.uint8).reshape(-1, H, W, 3)
small = frames[:, ::8, ::8].astype(np.float32).mean(-1)
diff = np.abs(np.diff(small, axis=0)).mean((1, 2))
cuts = np.where(diff > max(18.0, np.median(diff) * 6))[0] + 1 # first frame of each new shot
np.save(out / f"{video.stem}.npy", frames)
index.append({"name": video.stem, "frames": int(len(frames)), "width": W, "height": H, "cuts": cuts.tolist()})
print(video.stem, frames.shape, "cuts:", len(cuts))
(out / "index.json").write_text(json.dumps(index, indent=1))
|