Upload training/scripts/assign_buckets.py with huggingface_hub

training/scripts/assign_buckets.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""Stratified multi-resolution bucket assignment for normalized clips.
+
+Creates non-destructive symlink views (small/medium resolutions per user directive):
+  train_quick/v288 -> [512,288]  ~55%
+  train_quick/v384 -> [640,384]  ~28%
+  train_quick/v480 -> [832,480]  ~17%
+
+Constraint: a clip is only eligible for a bucket whose target area it can
+cover without upscaling. Deterministic (seed 42). Writes a manifest CSV.
+"""
+import csv
+import fractions
+import json
+import random
+import subprocess
+import sys
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+
+SRC = Path("/workspace/datasets/mega_24_ext")
+DST = Path("/workspace/datasets/train_quick")
+MANIFEST = Path("/workspace/projects/h3_loop/configs/bucket_manifest_quick.csv")
+
+BUCKETS = {  # name -> (w, h, share)
+    "v288": (512, 288, 0.55),
+    "v384": (640, 384, 0.28),
+    "v480": (832, 480, 0.17),
+}
+
+
+def probe(path: Path):
+    out = subprocess.run(
+        ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
+         "stream=width,height,nb_frames", "-of", "json", str(path)],
+        capture_output=True, text=True, timeout=120)
+    st = json.loads(out.stdout)["streams"][0]
+    return path.name, int(st["width"]), int(st["height"]), int(st.get("nb_frames") or 0)
+
+
+def main():
+    videos = sorted(SRC.glob("*.mp4"))
+    print(f"{len(videos)} normalized videos")
+    with ProcessPoolExecutor(max_workers=30) as ex:
+        info = list(ex.map(probe, videos, chunksize=16))
+
+    rng = random.Random(42)
+    areas = {name: w * h for name, (w, h, _) in BUCKETS.items()}
+    total = len(info)
+    quota = {name: int(total * share) for name, (_, _, share) in BUCKETS.items()}
+    # fill rounding remainder into v384
+    quota["v288"] += total - sum(quota.values())
+
+    # group by source resolution
+    groups = {}
+    for name, w, h, nb in info:
+        groups.setdefault((w, h), []).append(name)
+    for g in groups.values():
+        rng.shuffle(g)
+
+    assign = {}
+    counts = {b: 0 for b in BUCKETS}
+
+    def eligible(w, h):
+        return [b for b in BUCKETS if w * h >= areas[b]]
+
+    # 1) clips eligible for only one bucket go there
+    flexible = []
+    for (w, h), names in sorted(groups.items()):
+        el = eligible(w, h)
+        if not el:
+            for n in names:
+                assign[n] = "v288"  # tiny source: still train, bucket_no_upscale protects it
+                counts["v288"] += 1
+        elif len(el) == 1:
+            for n in names:
+                assign[n] = el[0]
+                counts[el[0]] += 1
+        else:
+            flexible.extend((n, w, h) for n in names)
+
+    # 2) flexible clips: fill scarcest-eligibility buckets first (v576 needs high-res)
+    rng.shuffle(flexible)
+    for bucket in ("v480", "v384", "v288"):
+        need = quota[bucket] - counts[bucket]
+        if need <= 0:
+            continue
+        rest = []
+        for n, w, h in flexible:
+            if need > 0 and bucket in eligible(w, h) and n not in assign:
+                assign[n] = bucket
+                counts[bucket] += 1
+                need -= 1
+            else:
+                rest.append((n, w, h))
+        flexible = rest
+    # leftovers -> largest eligible bucket below its quota, else v384
+    for n, w, h in flexible:
+        if n in assign:
+            continue
+        el = eligible(w, h)
+        assign[n] = el[-1] if el else "v288"
+        counts[assign[n]] += 1
+
+    print("assignment:", counts)
+
+    for b in BUCKETS:
+        d = DST / b
+        d.mkdir(parents=True, exist_ok=True)
+    made = 0
+    for name, bucket in assign.items():
+        src_mp4 = SRC / name
+        src_txt = src_mp4.with_suffix(".txt")
+        for s in (src_mp4, src_txt):
+            link = DST / bucket / s.name
+            if not link.exists() and s.exists():
+                link.symlink_to(s)
+                made += 1
+    print(f"{made} symlinks created")
+
+    nbmap = {name: (w, h, nb) for name, w, h, nb in info}
+    with open(MANIFEST, "w", newline="") as f:
+        wr = csv.writer(f)
+        wr.writerow(["file", "bucket", "src_width", "src_height", "src_frames_24fps"])
+        for name in sorted(assign):
+            w, h, nb = nbmap[name]
+            wr.writerow([name, assign[name], w, h, nb])
+    print(f"manifest: {MANIFEST}")
+
+
+if __name__ == "__main__":
+    main()