diff --git a/community_v3_migration/fix_dataset.py b/community_v3_migration/fix_dataset.py index af8978f6a..89132fb43 100644 --- a/community_v3_migration/fix_dataset.py +++ b/community_v3_migration/fix_dataset.py @@ -78,6 +78,21 @@ def _write_jsonl(path: Path, rows: list[dict]) -> None: f.write(json.dumps(r) + "\n") +def data_video_episode_mismatch(root) -> str | None: + """Return a description when the data files and any camera's video files disagree on the + episode count (dataset can't be migrated, e.g. 'All cams dont have same number of episodes'), + else None. Datasets without videos never mismatch here.""" + root = Path(root) + info = json.loads((root / "meta" / "info.json").read_text()) + counts = {"data": len(list((root / "data").glob("*/episode_*.parquet")))} + for k, f in info.get("features", {}).items(): + if f.get("dtype") == "video": + counts[k] = len(list((root / "videos").glob(f"*/{k}/episode_*.mp4"))) + if len(counts) > 1 and len(set(counts.values())) > 1: + return f"data/video episode counts disagree: {counts}" + return None + + def reconcile_episode_count(root) -> str | None: """When the data files and video files agree on an episode count N but the metadata lists a different count, rewrite the metadata (episodes.jsonl, episodes_stats.jsonl, info.json) to N. diff --git a/community_v3_migration/run_migration.py b/community_v3_migration/run_migration.py index 6a308ffb3..695b8b25c 100644 --- a/community_v3_migration/run_migration.py +++ b/community_v3_migration/run_migration.py @@ -17,7 +17,7 @@ from huggingface_hub import HfApi import so_arm_frame from classify import classify, is_end_effector, load_info -from fix_dataset import fix_dataset_in_place, reconcile_episode_count +from fix_dataset import data_video_episode_mismatch, fix_dataset_in_place, reconcile_episode_count SRC_REPO = "HuggingFaceVLA/community_dataset_v3" @@ -168,6 +168,10 @@ def migrate_one(api, dst_repo, sub, work_dir, no_upload) -> dict: if is_end_effector(info): return {"root": sub, "robot_type": info.get("robot_type"), "action": "skipped: end-effector (task-space) dataset, out of scope"} + mismatch = data_video_episode_mismatch(local) + if mismatch: + return {"root": sub, "robot_type": info.get("robot_type"), + "action": f"skipped: {mismatch}"} result = fix_dataset_in_place(local) # SO-arm value fix (or structural_only)