#!/usr/bin/env python3 """Generate data.json for the G1 Moves showcase site. Reads manifest.json + scans the filesystem to determine which media files (GIF/MP4) exist for each clip at each pipeline stage. Usage: python generate_data.py # from space/ directory python generate_data.py --repo /path/to/g1-moves """ import argparse import json import re from pathlib import Path STAGES = ["capture", "retarget", "training", "policy"] FEATURED_CLIPS = {"V_PullOver": 0} # Descriptive subtitles for karate moves (from README) KARATE_SUBTITLES = { "M_Move1": "Guard Combo", "M_Move2": "Low Punch", "M_Move3": "Horse Stance", "M_Move4": "Spin Punch", "M_Move5": "Twist Punch", "M_Move6": "Spin Strike", "M_Move7": "Rapid Punch", "M_Move8": "Drop Spin", "M_Move9": "Level Change", "M_Move10": "Side Kick", "M_Move11": "Blitz", "M_Move17": "Double Strike", "M_Move18": "Front Kick", "M_Move19": "Slow Kata", "M_Move20": "Open Strike", "M_ShortMove12": "Quick Jab", "M_ShortMove13": "Snap Kick", "M_ShortMove14": "Light Punch", "M_ShortMove15": "Drop Strike", "M_ShortMove16": "Power Burst", } def format_display_name(clip_id: str) -> str: """Convert clip ID to human-readable name.""" name = re.sub(r"^[BJMV]_", "", clip_id) name = re.sub(r"([a-z\d])([A-Z])", r"\1 \2", name) name = re.sub(r"([A-Za-z])(\d)", r"\1 \2", name) name = name.replace("_", " ") return name def get_media_file(clip_id: str, stage: str) -> str: """Return the expected GIF filename for a clip at a given stage.""" if stage == "capture": return f"{clip_id}.gif" return f"{clip_id}_{stage}.gif" def scan_stage(repo: Path, category: str, clip_id: str, stage: str) -> dict | None: """Check for media or pipeline artifacts at a pipeline stage. Visual GIF/MP4 previews are preferred. For newly added robot-compatible clips such as V_PullOver, some stages only have the real artifacts (CSV/PKL/NPZ/ONNX/PT) and no rendered preview yet. Return artifact metadata so the public page shows the stage as available instead of falsely saying "Not yet available". """ stage_dir = repo / category / clip_id / stage result = {} gif_name = get_media_file(clip_id, stage) gif_path = stage_dir / gif_name if gif_path.exists(): result["gif"] = f"{category}/{clip_id}/{stage}/{gif_name}" # Also check for MP4 mp4_name = gif_name.replace(".gif", ".mp4") mp4_path = stage_dir / mp4_name if mp4_path.exists(): result["mp4"] = f"{category}/{clip_id}/{stage}/{mp4_name}" artifacts = [] if stage == "capture": # The manifest entry itself means the source clip was captured/imported, # even if the original phone/video file is not published in the dataset. result.setdefault("status", "Source captured/imported") elif stage == "retarget": for suffix, label in [(".csv", "Retarget CSV"), (".pkl", "Robot motion PKL")]: path = stage_dir / f"{clip_id}{suffix}" if path.exists(): artifacts.append({"label": label, "url": f"{category}/{clip_id}/{stage}/{clip_id}{suffix}"}) elif stage == "training": path = stage_dir / f"{clip_id}.npz" if path.exists(): artifacts.append({"label": "Training NPZ", "url": f"{category}/{clip_id}/{stage}/{clip_id}.npz"}) elif stage == "policy": for policy_dir, label_prefix in [("policy", "Browser policy"), ("policy_154", "Robot-compatible 154D policy")]: for suffix, label in [(".onnx", "ONNX"), (".pt", "PyTorch checkpoint")]: path = repo / category / clip_id / policy_dir / f"{clip_id}_policy{suffix}" if path.exists(): artifacts.append({"label": f"{label_prefix} {label}", "url": f"{category}/{clip_id}/{policy_dir}/{clip_id}_policy{suffix}"}) if artifacts: result["artifacts"] = artifacts result.setdefault("status", "Artifacts available") return result if result else None def main(): parser = argparse.ArgumentParser(description="Generate site data from manifest") parser.add_argument( "--repo", type=Path, default=Path(__file__).parent.parent, help="Path to g1-moves repository root", ) args = parser.parse_args() repo = args.repo.resolve() manifest_path = repo / "manifest.json" if not manifest_path.exists(): print(f"ERROR: manifest.json not found at {manifest_path}") return manifest = json.loads(manifest_path.read_text()) clips = [] stats = {"dance": 0, "karate": 0, "bonus": 0, "policies": 0, "total": 0} for clip_id, info in sorted(manifest["clips"].items()): cat = info["category"] stats[cat] = stats.get(cat, 0) + 1 stats["total"] += 1 stages = {} for stage in STAGES: media = scan_stage(repo, cat, clip_id, stage) if media: stages[stage] = media # Check for browser-viewer policy files (legacy 160-dim) and robot policy_154 files. # policy_154 files are real-robot artifacts, but the current browser viewer still # expects the older media///policy/_policy.onnx path. onnx_path = repo / cat / clip_id / "policy" / f"{clip_id}_policy.onnx" pt_path = repo / cat / clip_id / "policy" / f"{clip_id}_policy.pt" policy154_onnx_path = repo / cat / clip_id / "policy_154" / f"{clip_id}_policy.onnx" policy154_pt_path = repo / cat / clip_id / "policy_154" / f"{clip_id}_policy.pt" has_onnx = onnx_path.exists() has_pt = pt_path.exists() has_policy_154 = policy154_onnx_path.exists() or policy154_pt_path.exists() has_policy = "policy" in stages or has_onnx or has_pt or has_policy_154 supports_viewer = has_onnx or has_policy_154 viewer_policy = "policy" if has_onnx else ("policy_154" if has_policy_154 else None) if has_policy: stats["policies"] += 1 subtitle = KARATE_SUBTITLES.get(clip_id) display_name = format_display_name(clip_id) if subtitle: display_name = f"{display_name}: {subtitle}" clips.append( { "id": clip_id, "name": display_name, "category": cat, "performer": info.get("performer", "Unknown"), "duration": info.get("duration_s", 0), "fps": info.get("fps", 60), "frames": info.get("frames", 0), "stages": stages, "has_policy": has_policy, "has_onnx": has_onnx, "has_policy_154": has_policy_154, "supports_viewer": supports_viewer, "viewer_policy": viewer_policy, "featured_rank": FEATURED_CLIPS.get(clip_id), } ) # Sort featured clips first, then categories grouped, then alphabetical. cat_order = {"dance": 0, "karate": 1, "bonus": 2} clips.sort(key=lambda c: ( c.get("featured_rank") if c.get("featured_rank") is not None else 999, cat_order.get(c["category"], 9), c["id"], )) data = { "base_url": "https://huggingface.co/datasets/exptech/g1-moves/resolve/main", "clips": clips, "stats": stats, } output = Path(__file__).parent / "data.json" output.write_text(json.dumps(data, indent=2)) print(f"Generated {output.name}: {len(clips)} clips, {stats['policies']} policies") if __name__ == "__main__": main()