| """Fetch + extract the WISA-80K wan480p shards. |
| |
| The clips live in ``csusupergear/WISA-80K-wan480p-16fps-81f`` as 276 webdataset |
| tars; shard ``N`` holds ids ``N*100 .. N*100+99`` as ``<id>.{mp4,txt,json}``. |
| Download and extraction are interleaved over a thread pool so a shard is unpacked |
| as soon as it lands. Both steps are idempotent: an already-downloaded shard is |
| served from the HF cache and an already-extracted shard is skipped. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| import tarfile |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from pathlib import Path |
|
|
| REPO = "csusupergear/WISA-80K-wan480p-16fps-81f" |
|
|
|
|
| def shard_done(vdir: Path, shard: int) -> bool: |
| """A shard counts as extracted when its last clip's three files exist.""" |
| last = f"{shard * 100 + 99:08d}" |
| return all((vdir / f"{last}.{ext}").exists() for ext in ("mp4", "txt", "json")) |
|
|
|
|
| def fetch_one(shard: int, raw_dir: Path, vdir: Path) -> tuple[int, str]: |
| from huggingface_hub import hf_hub_download |
|
|
| if shard_done(vdir, shard): |
| return shard, "skip" |
| path = hf_hub_download( |
| REPO, f"{shard:06d}.tar", repo_type="dataset", |
| local_dir=str(raw_dir), |
| ) |
| with tarfile.open(path) as tf: |
| tf.extractall(vdir) |
| return shard, "ok" |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--num_shards", type=int, default=276) |
| ap.add_argument("--raw_dir", default="data/raw/wisa_wan480p") |
| ap.add_argument("--video_dir", default="data/videos") |
| ap.add_argument("--workers", type=int, default=8) |
| args = ap.parse_args() |
|
|
| raw_dir, vdir = Path(args.raw_dir), Path(args.video_dir) |
| vdir.mkdir(parents=True, exist_ok=True) |
|
|
| done = n_err = 0 |
| with ThreadPoolExecutor(args.workers) as ex: |
| futs = {ex.submit(fetch_one, s, raw_dir, vdir): s |
| for s in range(args.num_shards)} |
| for f in as_completed(futs): |
| s = futs[f] |
| try: |
| _, status = f.result() |
| except Exception as e: |
| n_err += 1 |
| print(f"[err] shard {s}: {e}", flush=True) |
| continue |
| done += 1 |
| if done % 10 == 0: |
| print(f"{done}/{args.num_shards} shards ({status} {s:06d})", |
| flush=True) |
|
|
| n = len(list(vdir.glob("*.mp4"))) |
| print(f"DONE shards={done} errors={n_err} clips={n}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
