| |
| """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 = { |
| "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()} |
| |
| quota["v288"] += total - sum(quota.values()) |
|
|
| |
| 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]] |
|
|
| |
| flexible = [] |
| for (w, h), names in sorted(groups.items()): |
| el = eligible(w, h) |
| if not el: |
| for n in names: |
| assign[n] = "v288" |
| 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) |
|
|
| |
| 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 |
| |
| 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() |
|
|
