File size: 4,326 Bytes
cb7dd5b | 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 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | #!/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()
|