mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
migration: reconcile stale meta episode count to data+video files
When the data files and video files agree on episode count N but the metadata lists a different count, rewrite meta/episodes.jsonl, meta/episodes_stats.jsonl and info.json to N before the v2.1->v3.0 converter runs (which otherwise raises 'Number of episodes is not the same'). Only the safe direction is handled: trimming metadata that lists MORE episodes than exist. Non-contiguous data, data/video disagreement, or metadata missing episodes are left untouched.
This commit is contained in:
@@ -67,6 +67,61 @@ def _regen_episode_stats(root: Path) -> None:
|
|||||||
f.write(json.dumps(orig[ep]) + "\n")
|
f.write(json.dumps(orig[ep]) + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
def _read_jsonl(path: Path) -> list[dict]:
|
||||||
|
with open(path) as f:
|
||||||
|
return [json.loads(line) for line in f if line.strip()]
|
||||||
|
|
||||||
|
|
||||||
|
def _write_jsonl(path: Path, rows: list[dict]) -> None:
|
||||||
|
with open(path, "w") as f:
|
||||||
|
for r in rows:
|
||||||
|
f.write(json.dumps(r) + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
Only the safe direction is handled: trimming metadata that lists MORE episodes than actually
|
||||||
|
exist. If the data itself is non-contiguous, the videos disagree with the data, or the metadata
|
||||||
|
lists FEWER episodes than the data (which would require fabricating per-episode stats), nothing
|
||||||
|
is changed and the stock converter's mismatch error is left to surface. Returns a note on fix."""
|
||||||
|
root = Path(root)
|
||||||
|
info = json.loads((root / "meta" / "info.json").read_text())
|
||||||
|
|
||||||
|
data_idx = sorted(int(p.stem.split("_")[-1]) for p in (root / "data").glob("*/episode_*.parquet"))
|
||||||
|
n = len(data_idx)
|
||||||
|
if n == 0 or data_idx != list(range(n)):
|
||||||
|
return None
|
||||||
|
|
||||||
|
vkeys = [k for k, f in info.get("features", {}).items() if f.get("dtype") == "video"]
|
||||||
|
for k in vkeys:
|
||||||
|
if len(list((root / "videos").glob(f"*/{k}/episode_*.mp4"))) != n:
|
||||||
|
return None # data and videos disagree -> out of scope for this fix
|
||||||
|
|
||||||
|
eps_path = root / "meta" / "episodes.jsonl"
|
||||||
|
stats_path = root / "meta" / "episodes_stats.jsonl"
|
||||||
|
eps, stats = _read_jsonl(eps_path), _read_jsonl(stats_path)
|
||||||
|
if len(eps) == n and len(stats) == n:
|
||||||
|
return None
|
||||||
|
|
||||||
|
eps_keep = [e for e in eps if e.get("episode_index", -1) < n]
|
||||||
|
stats_keep = [s for s in stats if s.get("episode_index", -1) < n]
|
||||||
|
if len(eps_keep) != n or len(stats_keep) != n:
|
||||||
|
return None # metadata is missing episodes present in the data -> can't safely fabricate
|
||||||
|
|
||||||
|
dropped = max(len(eps), len(stats)) - n
|
||||||
|
_write_jsonl(eps_path, eps_keep)
|
||||||
|
_write_jsonl(stats_path, stats_keep)
|
||||||
|
info["total_episodes"] = n
|
||||||
|
info["total_frames"] = int(sum(e.get("length", 0) for e in eps_keep))
|
||||||
|
if "total_videos" in info:
|
||||||
|
info["total_videos"] = n * len(vkeys)
|
||||||
|
info["splits"] = {"train": f"0:{n}"}
|
||||||
|
(root / "meta" / "info.json").write_text(json.dumps(info, indent=4))
|
||||||
|
return f"metadata episode count reconciled to {n} (data & videos agree; dropped {dropped} stale meta entries)"
|
||||||
|
|
||||||
|
|
||||||
def fix_dataset_in_place(root) -> dict:
|
def fix_dataset_in_place(root) -> dict:
|
||||||
"""Returns the classification dict augmented with the action taken."""
|
"""Returns the classification dict augmented with the action taken."""
|
||||||
root = Path(root)
|
root = Path(root)
|
||||||
|
|||||||
@@ -17,7 +17,7 @@ from huggingface_hub import HfApi
|
|||||||
|
|
||||||
import so_arm_frame
|
import so_arm_frame
|
||||||
from classify import classify, is_end_effector, load_info
|
from classify import classify, is_end_effector, load_info
|
||||||
from fix_dataset import fix_dataset_in_place
|
from fix_dataset import fix_dataset_in_place, reconcile_episode_count
|
||||||
|
|
||||||
SRC_REPO = "HuggingFaceVLA/community_dataset_v3"
|
SRC_REPO = "HuggingFaceVLA/community_dataset_v3"
|
||||||
|
|
||||||
@@ -171,6 +171,10 @@ def migrate_one(api, dst_repo, sub, work_dir, no_upload) -> dict:
|
|||||||
|
|
||||||
result = fix_dataset_in_place(local) # SO-arm value fix (or structural_only)
|
result = fix_dataset_in_place(local) # SO-arm value fix (or structural_only)
|
||||||
|
|
||||||
|
reconciled = reconcile_episode_count(local) # align stale meta counts to data+video files
|
||||||
|
if reconciled:
|
||||||
|
result["action"] = f"{result['action']}; {reconciled}"
|
||||||
|
|
||||||
from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset
|
from lerobot.scripts.convert_dataset_v21_to_v30 import convert_dataset
|
||||||
convert_dataset(repo_id=sub, root=str(local), push_to_hub=False) # v2.1 -> v3.0, in place
|
convert_dataset(repo_id=sub, root=str(local), push_to_hub=False) # v2.1 -> v3.0, in place
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user