processing/build_depth_dataset.py

#!/usr/bin/env python3
"""Randomly select videos present in the caption file, the selected videos
dir, and depth/, then materialize them into a flat depth-control dataset
directory (mirroring the layout built by
HumanVID/processing/build_pose_dataset.py): each selected video becomes a
{stem}.mp4 / {stem}_depth.mp4 pair plus a dataset.json describing
media_path, reference_path, and caption.
"""

import argparse
import json
import os
import random
import shutil
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

_HERE = Path(__file__).resolve().parent
_VIDGEN_ROOT = _HERE.parent


def load_captions(caption_file: Path) -> dict[str, str]:
    """Load a {vid: caption} lookup from a VidGen-style caption JSON file
    (a list of {"vid": ..., "caption": ...} entries)."""
    entries = json.loads(caption_file.read_text(encoding="utf-8"))
    return {
        entry["vid"]: entry["caption"]
        for entry in entries
        if entry.get("vid") is not None and entry.get("caption") is not None
    }


def find_eligible_stems(videos_dir: Path, captions: dict, depth_dir: Path) -> list[str]:
    """Return stems present in videos_dir, captions, and depth_dir."""
    eligible = []
    for video_path in sorted(videos_dir.glob("*.mp4")):
        stem = video_path.stem
        if stem not in captions:
            continue
        if not (depth_dir / f"{stem}_depth.mp4").exists():
            continue
        eligible.append(stem)
    return eligible


def link_or_copy(src: Path, dst: Path, mode: str) -> None:
    dst.unlink(missing_ok=True)
    if mode == "symlink":
        dst.symlink_to(src.resolve())
    elif mode == "hardlink":
        try:
            os.link(src, dst)
        except OSError:
            shutil.copy2(src, dst)
    else:
        shutil.copy2(src, dst)


def materialize_pair(stem: str, media_src: Path, reference_src: Path, output_dir: Path, link_mode: str):
    try:
        link_or_copy(media_src, output_dir / f"{stem}.mp4", link_mode)
        link_or_copy(reference_src, output_dir / f"{stem}_depth.mp4", link_mode)
        return stem, True, None
    except OSError as e:
        return stem, False, str(e)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--videos-dir", type=Path,
                         default=_VIDGEN_ROOT / "selected_videos",
                         help="Directory containing the selected videos")
    parser.add_argument("--captions", type=Path,
                         default=_VIDGEN_ROOT / "VidGen_1M_video_caption.json",
                         help="Path to a VidGen-style caption JSON (list of {vid, caption})")
    parser.add_argument("--depth-dir", type=Path, default=None,
                         help="Directory of {stem}_depth.mp4 files (default: <videos-dir>/depth)")
    parser.add_argument("--output", type=Path, required=True,
                         help="Output dataset directory to create")
    parser.add_argument("--num-videos", type=int, default=2000,
                         help="Number of videos to randomly select (default: 2000)")
    parser.add_argument("--seed", type=int, default=0,
                         help="Random seed for reproducible selection")
    parser.add_argument("--link-mode", choices=["copy", "symlink", "hardlink"], default="hardlink",
                         help="How to materialize files into --output (default: hardlink; "
                              "falls back to a real copy if hardlinking fails, e.g. across filesystems)")
    parser.add_argument("--workers", type=int, default=16,
                         help="Number of concurrent copy/link workers")
    args = parser.parse_args()

    depth_dir = args.depth_dir or args.videos_dir / "depth"

    if not args.videos_dir.is_dir():
        sys.exit(f"videos dir not found: {args.videos_dir}")
    if not args.captions.is_file():
        sys.exit(f"captions file not found: {args.captions}")
    if not depth_dir.is_dir():
        sys.exit(f"depth dir not found: {depth_dir}")

    captions = load_captions(args.captions)

    eligible = find_eligible_stems(args.videos_dir, captions, depth_dir)
    print(f"Found {len(eligible)} videos present in {args.captions.name}, {args.videos_dir.name}/, "
          f"and {depth_dir.name}/", file=sys.stderr)

    if len(eligible) < args.num_videos:
        print(f"Warning: only {len(eligible)} eligible videos available, "
              f"fewer than requested {args.num_videos}. Using all of them.", file=sys.stderr)
        selected = eligible
    else:
        rng = random.Random(args.seed)
        selected = rng.sample(eligible, args.num_videos)
    selected.sort()

    args.output.mkdir(parents=True, exist_ok=True)

    results = {}
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = [
            pool.submit(
                materialize_pair,
                stem,
                args.videos_dir / f"{stem}.mp4",
                depth_dir / f"{stem}_depth.mp4",
                args.output,
                args.link_mode,
            )
            for stem in selected
        ]
        done = 0
        for fut in as_completed(futures):
            stem, ok, err = fut.result()
            results[stem] = (ok, err)
            done += 1
            print(f"\r{done}/{len(selected)} materialized", end="", file=sys.stderr)
    print(file=sys.stderr)

    dataset = []
    failed = []
    for stem in selected:
        ok, err = results[stem]
        if not ok:
            failed.append((stem, err))
            continue
        dataset.append({
            "caption": captions[stem],
            "media_path": f"{stem}.mp4",
            "reference_path": f"{stem}_depth.mp4",
        })

    if failed:
        print(f"Warning: {len(failed)} video(s) failed to materialize:", file=sys.stderr)
        for stem, err in failed:
            print(f"  {stem}: {err}", file=sys.stderr)

    dataset_json_path = args.output / "dataset.json"
    dataset_json_path.write_text(json.dumps(dataset, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

    print(f"Done. Wrote {len(dataset)} pairs to {args.output} "
          f"(reference source: {depth_dir}/, link mode: {args.link_mode}), "
          f"manifest at {dataset_json_path}")


if __name__ == "__main__":
    main()