training/scripts/assign_buckets.py

#!/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()