#!/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}")