File size: 2,160 Bytes
b152c62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | #!/usr/bin/env python
"""Enumerate unitreerobotics LeRobot datasets and emit a conversion worklist.
Filters to repos whose meta/info.json has flat observation.state + action
features (the v2.1 G1_Dex1_* sets split state into 7 named fields — those need
a concat converter and are excluded here). Output: JSON lines sorted by frames,
ready to drive build_shards(.py|_v2.py) + pack_wds.py.
Usage:
python list_unitree_specs.py [--min-frames 10000] > unitree_worklist.jsonl
"""
from __future__ import annotations
import argparse
import json
import sys
import requests
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--org", default="unitreerobotics")
ap.add_argument("--min-frames", type=int, default=0)
args = ap.parse_args()
ds = requests.get(
f"https://huggingface.co/api/datasets?author={args.org}&limit=200", timeout=30
).json()
rows = []
for d in ds:
rid = d["id"]
try:
info = requests.get(
f"https://huggingface.co/datasets/{rid}/resolve/main/meta/info.json",
timeout=30, allow_redirects=True,
).json()
except Exception:
continue
feats = info.get("features", {})
if "observation.state" not in feats or "action" not in feats:
print(f"skip {rid}: no flat state/action", file=sys.stderr)
continue
cams = [k for k, v in feats.items() if v.get("dtype") == "video"]
rows.append({
"repo_id": rid,
"version": info["codebase_version"],
"fps": info["fps"],
"episodes": info["total_episodes"],
"frames": info["total_frames"],
"state_dim": feats["observation.state"]["shape"][0],
"action_dim": feats["action"]["shape"][0],
"cameras": cams,
})
rows.sort(key=lambda r: -r["frames"])
kept = [r for r in rows if r["frames"] >= args.min_frames]
for r in kept:
print(json.dumps(r))
tot = sum(r["frames"] for r in kept)
print(f"{len(kept)} datasets, {tot:,} frames total", file=sys.stderr)
if __name__ == "__main__":
main()
|