migration: ditch datasets whose data and camera files disagree on episode count

If the data files and any camera's video files (or two cameras) don't have the
same number of episodes, skip the dataset up front instead of letting the
converter raise 'All cams dont have same number of episodes' mid-run.
This commit is contained in:
CarolinePascal
2026-07-17 17:00:29 +02:00
parent a6dcb18585
commit e0d455ec9f
2 changed files with 20 additions and 1 deletions
+15
View File
@@ -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.
+5 -1
View File
@@ -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)