generate_data.py

#!/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/<category>/<clip>/policy/<clip>_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()