Upload training/scripts/normalize_videos.py with huggingface_hub

training/scripts/normalize_videos.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Normalize mega_curated videos: 24 fps CFR, square pixels (SAR 1:1), copy captions.
+
+Non-destructive: writes derived files to OUT_DIR, sources untouched.
+Skips clips whose 24fps-converted frame count would fall below 22 (17n+5 grid minimum for video).
+Usage: normalize_videos.py SRC_DIR OUT_DIR [--jobs N]
+"""
+import argparse
+import fractions
+import json
+import shutil
+import subprocess
+import sys
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from pathlib import Path
+
+MIN_FRAMES = 22  # smallest useful 17n+5 video length
+
+
+def probe(path: Path):
+    out = subprocess.run(
+        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
+         "stream=width,height,sample_aspect_ratio,avg_frame_rate,nb_frames,duration",
+         "-of", "json", str(path)],
+        capture_output=True, text=True, timeout=120)
+    if out.returncode != 0:
+        return None
+    st = json.loads(out.stdout)["streams"][0]
+    w, h = int(st["width"]), int(st["height"])
+    sar = st.get("sample_aspect_ratio") or "1:1"
+    if sar in ("0:1", "N/A"):
+        sar = "1:1"
+    fps = float(fractions.Fraction(st["avg_frame_rate"]))
+    nb = st.get("nb_frames")
+    if nb is None or nb == "N/A":
+        dur = float(st.get("duration") or 0)
+        nb = int(dur * fps)
+    return w, h, sar, fps, int(nb)
+
+
+def process(src: Path, dst_dir: Path):
+    info = probe(src)
+    if info is None:
+        return ("probe_fail", src.name, None)
+    w, h, sar, fps, nb = info
+    out_frames = int(nb * 24.0 / fps) if fps > 0 else 0
+    if out_frames < MIN_FRAMES:
+        return ("too_short", src.name, out_frames)
+    sn, sd = (int(x) for x in sar.split(":"))
+    vf = ["fps=24"]
+    if sn != sd:  # anamorphic: resample to square pixels, keep height
+        out_w = round(w * sn / sd / 2) * 2
+        vf.insert(0, f"scale={out_w}:{h}:flags=lanczos")
+        final_w = out_w
+    else:
+        final_w = w
+    vf.append("setsar=1")
+    dst = dst_dir / src.name
+    if dst.exists():
+        return ("exists", src.name, (final_w, h))
+    r = subprocess.run(
+        ["ffmpeg", "-y", "-v", "error", "-i", str(src),
+         "-vf", ",".join(vf), "-c:v", "libx264", "-crf", "15",
+         "-preset", "fast", "-pix_fmt", "yuv420p", "-an", "-map_metadata", "-1",
+         str(dst) + ".tmp.mp4"],
+        capture_output=True, text=True, timeout=600)
+    if r.returncode != 0:
+        Path(str(dst) + ".tmp.mp4").unlink(missing_ok=True)
+        return ("encode_fail", src.name, r.stderr[-200:])
+    Path(str(dst) + ".tmp.mp4").rename(dst)
+    txt = src.with_suffix(".txt")
+    if txt.exists():
+        shutil.copy2(txt, dst_dir / txt.name)
+    return ("ok", src.name, (final_w, h))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("src")
+    ap.add_argument("out")
+    ap.add_argument("--jobs", type=int, default=30)
+    args = ap.parse_args()
+    src_dir, out_dir = Path(args.src), Path(args.out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    videos = sorted(src_dir.glob("*.mp4"))
+    print(f"{len(videos)} source videos", flush=True)
+    results = {"ok": 0, "exists": 0, "too_short": 0, "probe_fail": 0, "encode_fail": 0}
+    failures = []
+    with ProcessPoolExecutor(max_workers=args.jobs) as ex:
+        futs = {ex.submit(process, v, out_dir): v for v in videos}
+        for i, fut in enumerate(as_completed(futs)):
+            status, name, extra = fut.result()
+            results[status] += 1
+            if status in ("too_short", "probe_fail", "encode_fail"):
+                failures.append((status, name, extra))
+            if (i + 1) % 500 == 0:
+                print(f"{i+1}/{len(videos)} {results}", flush=True)
+    print("FINAL:", results, flush=True)
+    with open(out_dir.parent / "normalize_failures.log", "w") as f:
+        for status, name, extra in failures:
+            f.write(f"{status}\t{name}\t{extra}\n")
+
+
+if __name__ == "__main__":
+    main()