build_index.py

#!/usr/bin/env python3
"""
Build a filtered training index from community_dataset_v3 on disk.

Applies:
  - Robot type filter (so100/so101 variants only)
  - Schema filter (2 cameras, 6-DOF, 30fps)
  - Episode length filter (5s-60s)
  - Per-task cap (default 200)
  - Per-contributor cap (default 200)
  - Excludes datasets with file count mismatches

Outputs filtered_index.json with all info needed to train.
"""

import argparse
import glob
import json
import random
from collections import defaultdict
from pathlib import Path

import av
import pandas as pd


def get_video_duration(video_path: Path) -> float:
    """Get video duration in seconds by reading container metadata (fast, no decoding)."""
    try:
        container = av.open(str(video_path))
        stream = container.streams.video[0]
        duration = float(stream.duration * stream.time_base)
        container.close()
        return duration
    except Exception:
        return 0.0


def load_dataset_meta(dataset_root: Path) -> dict | None:
    """Load and validate a single dataset's metadata."""
    info_path = dataset_root / "meta" / "info.json"
    if not info_path.exists():
        return None

    info = json.load(open(info_path))

    # Robot type filter
    robot = info.get("robot_type", "")
    if robot not in ("so100", "so101", "so100_follower", "so101_follower"):
        return None

    # Schema filter: exactly the 2-camera, 6-DOF schema
    features = info.get("features", {})
    expected_keys = {
        "action", "episode_index", "frame_index", "index",
        "observation.images.image", "observation.images.image2",
        "observation.state", "task_index", "timestamp",
    }
    if set(features.keys()) != expected_keys:
        return None

    # Dimension check
    if features.get("action", {}).get("shape") != [6]:
        return None
    if features.get("observation.state", {}).get("shape") != [6]:
        return None

    # FPS check
    if info.get("fps") != 30:
        return None

    # Resolution check
    for cam_key in ("observation.images.image", "observation.images.image2"):
        shape = features.get(cam_key, {}).get("shape", [])
        if len(shape) < 2 or shape[0] != 480 or shape[1] != 640:
            return None

    # Load tasks
    tasks_path = dataset_root / "meta" / "tasks.jsonl"
    tasks = {}
    if tasks_path.exists():
        for line in open(tasks_path):
            line = line.strip()
            if line:
                t = json.loads(line)
                tasks[t["task_index"]] = t["task"]

    # Integrity check: video and parquet file counts
    total_eps = info.get("total_episodes", 0)
    vids = glob.glob(str(dataset_root / "videos" / "**" / "*.mp4"), recursive=True)
    parquets = glob.glob(str(dataset_root / "data" / "**" / "*.parquet"), recursive=True)
    expected_vids = total_eps * 2  # 2 cameras
    if len(vids) != expected_vids or len(parquets) != total_eps:
        return None

    # Load episode metadata if available
    episodes = []
    ep_jsonl = dataset_root / "meta" / "episodes.jsonl"
    if ep_jsonl.exists():
        for line in open(ep_jsonl):
            line = line.strip()
            if line:
                episodes.append(json.loads(line))

    return {
        "robot_type": robot,
        "total_episodes": total_eps,
        "total_frames": info.get("total_frames", 0),
        "fps": info["fps"],
        "tasks": tasks,
        "episodes": episodes,
        "features": {k: v.get("shape") for k, v in features.items()},
    }


def build_index(
    data_root: Path,
    max_per_task: int = 200,
    max_per_contributor: int = 200,
    min_episode_frames: int = 150,
    max_episode_frames: int = 1800,
    seed: int = 42,
) -> dict:
    """Build filtered training index."""
    rng = random.Random(seed)

    # Discover all contributor/dataset pairs
    contributors = sorted([
        d for d in data_root.iterdir()
        if d.is_dir() and not d.name.startswith(".")
    ])

    # Phase 1: Load all valid datasets
    all_episodes = []  # (contributor, dataset_name, episode_idx, task, num_frames)
    datasets_passed = 0
    datasets_rejected = 0
    skipped_missing = 0
    skipped_video_mismatch = 0

    for contrib_dir in contributors:
        if not contrib_dir.is_dir():
            continue
        contributor = contrib_dir.name

        for ds_dir in sorted(contrib_dir.iterdir()):
            if not ds_dir.is_dir():
                continue

            meta = load_dataset_meta(ds_dir)
            if meta is None:
                datasets_rejected += 1
                continue

            datasets_passed += 1
            dataset_name = f"{contributor}/{ds_dir.name}"

            # Default task if none specified
            if not meta["tasks"]:
                meta["tasks"] = {0: "(no task)"}

            # Build episode list by reading actual parquet files
            # Trust the parquet row count, not metadata
            for ep_idx in range(meta["total_episodes"]):
                parquet_path = ds_dir / f"data/chunk-000/episode_{ep_idx:06d}.parquet"
                if not parquet_path.exists():
                    skipped_missing += 1
                    continue

                # Read actual row count and timestamps from parquet
                pf_full = pd.read_parquet(parquet_path, columns=["frame_index", "timestamp"])
                actual_length = len(pf_full)

                if actual_length < min_episode_frames or actual_length > max_episode_frames:
                    continue

                # Also verify both video files exist
                vid1 = ds_dir / f"videos/chunk-000/observation.images.image/episode_{ep_idx:06d}.mp4"
                vid2 = ds_dir / f"videos/chunk-000/observation.images.image2/episode_{ep_idx:06d}.mp4"
                if not vid1.exists() or not vid2.exists():
                    skipped_missing += 1
                    continue

                # Verify video duration covers all parquet timestamps
                # The last frame's timestamp must be within the video duration
                last_timestamp = float(pf_full["timestamp"].iloc[-1])
                vid1_duration = get_video_duration(vid1)
                vid2_duration = get_video_duration(vid2)
                min_vid_duration = min(vid1_duration, vid2_duration)
                if min_vid_duration > 0 and last_timestamp > min_vid_duration:
                    # Video is shorter than parquet claims — truncate to what the video covers
                    # Find the last frame index where timestamp <= video duration
                    valid_mask = pf_full["timestamp"] <= min_vid_duration
                    actual_length = int(valid_mask.sum())
                    if actual_length < min_episode_frames:
                        skipped_video_mismatch += 1
                        continue

                # Get task from episodes.jsonl if available, else default
                task_idx = 0
                if meta["episodes"]:
                    for ep_meta in meta["episodes"]:
                        if ep_meta.get("episode_index") == ep_idx:
                            task_idx = ep_meta.get("task_index", 0)
                            break

                task = meta["tasks"].get(task_idx, "(no task)")
                all_episodes.append((contributor, dataset_name, ep_idx, task, actual_length))

    print(f"Datasets: {datasets_passed} passed, {datasets_rejected} rejected")
    print(f"Episodes verified: {len(all_episodes)}, skipped missing: {skipped_missing}, skipped video mismatch: {skipped_video_mismatch}")
    print(f"Episodes before caps: {len(all_episodes)}")

    # Phase 2: Apply per-task cap
    task_buckets = defaultdict(list)
    for ep in all_episodes:
        task_buckets[ep[3]].append(ep)

    after_task_cap = []
    tasks_capped = 0
    for task, eps in task_buckets.items():
        rng.shuffle(eps)
        if len(eps) > max_per_task:
            tasks_capped += 1
        after_task_cap.extend(eps[:max_per_task])

    print(f"Episodes after per-task cap ({max_per_task}): {len(after_task_cap)} ({tasks_capped} tasks capped)")

    # Phase 3: Apply per-contributor cap
    contrib_buckets = defaultdict(list)
    for ep in after_task_cap:
        contrib_buckets[ep[0]].append(ep)

    final_episodes = []
    contribs_capped = 0
    for contributor, eps in contrib_buckets.items():
        rng.shuffle(eps)
        if len(eps) > max_per_contributor:
            contribs_capped += 1
        final_episodes.extend(eps[:max_per_contributor])

    print(f"Episodes after per-contributor cap ({max_per_contributor}): {len(final_episodes)} ({contribs_capped} contributors capped)")

    # Phase 4: Build the index
    # Sort for determinism
    final_episodes.sort(key=lambda x: (x[1], x[2]))

    # Collect unique tasks
    unique_tasks = sorted(set(ep[3] for ep in final_episodes))
    task_to_idx = {t: i for i, t in enumerate(unique_tasks)}

    # Collect unique datasets used
    datasets_used = sorted(set(ep[1] for ep in final_episodes))

    # Build episode entries
    entries = []
    total_frames = 0
    for contributor, dataset_name, ep_idx, task, num_frames in final_episodes:
        entries.append({
            "dataset": dataset_name,
            "episode_index": ep_idx,
            "task": task,
            "task_index": task_to_idx[task],
            "num_frames": num_frames,
        })
        total_frames += num_frames

    index = {
        "source_repo": "HuggingFaceVLA/community_dataset_v3",
        "filters": {
            "max_per_task": max_per_task,
            "max_per_contributor": max_per_contributor,
            "min_episode_frames": min_episode_frames,
            "max_episode_frames": max_episode_frames,
            "seed": seed,
        },
        "summary": {
            "datasets": len(datasets_used),
            "episodes": len(entries),
            "unique_tasks": len(unique_tasks),
            "total_frames": total_frames,
            "est_hours": total_frames / 30 / 3600,
        },
        "tasks": unique_tasks,
        "datasets_used": datasets_used,
        "episodes": entries,
    }

    return index


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-root", type=Path, default=Path.home() / "lap" / "community_dataset_v3")
    parser.add_argument("--output", type=Path, default=Path(__file__).parent / "filtered_index.json")
    parser.add_argument("--max-per-task", type=int, default=200)
    parser.add_argument("--max-per-contributor", type=int, default=200)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    index = build_index(
        args.data_root,
        max_per_task=args.max_per_task,
        max_per_contributor=args.max_per_contributor,
        seed=args.seed,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with open(args.output, "w") as f:
        json.dump(index, f, indent=2)

    print(f"\nSaved to {args.output}")
    print(f"  Datasets: {index['summary']['datasets']}")
    print(f"  Episodes: {index['summary']['episodes']}")
    print(f"  Tasks: {index['summary']['unique_tasks']}")
    print(f"  Frames: {index['summary']['total_frames']:,}")
    print(f"  Est. hours: {index['summary']['est_hours']:.1f}")