Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,268 @@
+#!/usr/bin/env python3
+"""
+DI LeRobot Converter API
+========================
+Receives episode data (JSON + video URL) from the iOS app,
+creates a LeRobot v2.0 parquet file, uploads parquet + video
+to the HuggingFace dataset repo, and updates meta/info.json.
+
+Deployed as a HuggingFace Space with Gradio.
+The iOS app calls the /api/convert endpoint after uploading to GCS.
+"""
+
+import gradio as gr
+import json
+import os
+import tempfile
+import shutil
+from pathlib import Path
+
+import pandas as pd
+import numpy as np
+from huggingface_hub import HfApi, hf_hub_download
+
+# Config
+HF_DATASET_REPO = "DynamicIntelligence/humanoid-robots-training-dataset"
+HF_TOKEN = os.environ.get("HF_TOKEN", "")
+GCS_BUCKET = "di_record_intern_data"
+CHUNKS_SIZE = 100
+
+
+def convert_episode(episode_json: str) -> str:
+    """
+    Convert episode data to LeRobot v2.0 format and upload to dataset repo.
+
+    Input JSON schema:
+    {
+        "episode_index": int,          # auto-assigned if -1
+        "language_instruction": str,
+        "fps": int,
+        "frames": [
+            {
+                "timestamp": float,
+                "pose": {"x": f, "y": f, "z": f, "yaw": f, "pitch": f, "roll": f},
+                "left_hand": [x, y, z] or null,
+                "right_hand": [x, y, z] or null
+            }, ...
+        ],
+        "video_gcs_path": str          # GCS path to rgb_video.mp4
+    }
+    """
+    try:
+        data = json.loads(episode_json)
+    except json.JSONDecodeError as e:
+        return json.dumps({"error": f"Invalid JSON: {e}"})
+
+    api = HfApi(token=HF_TOKEN)
+
+    # Determine episode index
+    episode_index = data.get("episode_index", -1)
+    if episode_index < 0:
+        # Auto-assign: read current info.json to get next index
+        try:
+            info_path = hf_hub_download(
+                repo_id=HF_DATASET_REPO, filename="meta/info.json",
+                repo_type="dataset", token=HF_TOKEN
+            )
+            with open(info_path) as f:
+                info = json.load(f)
+            episode_index = info.get("total_episodes", 0)
+        except Exception:
+            episode_index = 0
+
+    lang = data.get("language_instruction", "")
+    fps = data.get("fps", 30) or 30
+    frames = data.get("frames", [])
+    num_frames = len(frames)
+
+    if num_frames == 0:
+        return json.dumps({"error": "No frames in episode data"})
+
+    # Build parquet rows
+    rows = []
+    for i, frame in enumerate(frames):
+        pose = frame.get("pose", {})
+        cam_x = pose.get("x", 0)
+        cam_y = pose.get("y", 0)
+        cam_z = pose.get("z", 0)
+        cam_roll = pose.get("roll", 0)
+        cam_pitch = pose.get("pitch", 0)
+        cam_yaw = pose.get("yaw", 0)
+        camera_pose = [cam_x, cam_y, cam_z, cam_roll, cam_pitch, cam_yaw]
+
+        # Hand data: [x, y, z] from end_effector → pad to 9 values (3 joints × xyz)
+        lh = frame.get("left_hand") or [0, 0, 0]
+        rh = frame.get("right_hand") or [0, 0, 0]
+        # Pad single palm position to 3-joint format (wrist=palm, others=0)
+        left_hand = list(lh[:3]) + [0.0] * 6
+        right_hand = list(rh[:3]) + [0.0] * 6
+
+        # Action deltas
+        if i > 0:
+            prev = frames[i - 1]
+            pp = prev.get("pose", {})
+            prev_cam = [pp.get("x", 0), pp.get("y", 0), pp.get("z", 0),
+                        pp.get("roll", 0), pp.get("pitch", 0), pp.get("yaw", 0)]
+            cam_delta = [camera_pose[j] - prev_cam[j] for j in range(6)]
+
+            plh = prev.get("left_hand") or [0, 0, 0]
+            prh = prev.get("right_hand") or [0, 0, 0]
+            lh_delta = [lh[j] - plh[j] if j < len(lh) and j < len(plh) else 0 for j in range(3)] + [0.0] * 6
+            rh_delta = [rh[j] - prh[j] if j < len(rh) and j < len(prh) else 0 for j in range(3)] + [0.0] * 6
+        else:
+            cam_delta = [0.0] * 6
+            lh_delta = [0.0] * 9
+            rh_delta = [0.0] * 9
+
+        rows.append({
+            "episode_index": episode_index,
+            "frame_index": i,
+            "timestamp": frame.get("timestamp", i / fps),
+            "observation.camera_pose": camera_pose,
+            "observation.left_hand": left_hand,
+            "observation.right_hand": right_hand,
+            "action.camera_delta": cam_delta,
+            "action.left_hand_delta": lh_delta,
+            "action.right_hand_delta": rh_delta,
+            "language_instruction": lang,
+            "next.done": i == num_frames - 1,
+        })
+
+    # Create parquet
+    tmp = Path(tempfile.mkdtemp())
+    try:
+        df = pd.DataFrame(rows)
+        chunk_idx = episode_index // CHUNKS_SIZE
+        parquet_path = tmp / f"episode_{episode_index:06d}.parquet"
+        df.to_parquet(parquet_path, index=False)
+
+        # Upload parquet
+        api.upload_file(
+            path_or_fileobj=str(parquet_path),
+            path_in_repo=f"data/chunk-{chunk_idx:03d}/episode_{episode_index:06d}.parquet",
+            repo_id=HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN,
+        )
+
+        # Upload video from GCS if provided
+        video_gcs_path = data.get("video_gcs_path", "")
+        video_gcs_url = data.get("video_gcs_url", "")
+        video_uploaded = False
+
+        if video_gcs_url:
+            # Download from GCS public URL and re-upload to HF
+            import urllib.request
+            video_local = tmp / "rgb_video.mp4"
+            try:
+                urllib.request.urlretrieve(video_gcs_url, str(video_local))
+                api.upload_file(
+                    path_or_fileobj=str(video_local),
+                    path_in_repo=f"videos/chunk-{chunk_idx:03d}/rgb/episode_{episode_index:06d}.mp4",
+                    repo_id=HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN,
+                )
+                video_uploaded = True
+            except Exception as ve:
+                pass  # Video upload is optional
+
+        # Update meta/info.json
+        try:
+            existing_info_path = hf_hub_download(
+                repo_id=HF_DATASET_REPO, filename="meta/info.json",
+                repo_type="dataset", token=HF_TOKEN
+            )
+            with open(existing_info_path) as f:
+                info = json.load(f)
+            info["total_episodes"] = max(info.get("total_episodes", 0), episode_index + 1)
+            info["total_frames"] = info.get("total_frames", 0) + num_frames
+            info["splits"] = {"train": f"0:{info['total_episodes']}"}
+            info["total_chunks"] = (info["total_episodes"] - 1) // CHUNKS_SIZE + 1
+            if video_uploaded:
+                info["total_videos"] = info.get("total_videos", 0) + 1
+        except Exception:
+            info = build_default_info(episode_index, num_frames)
+
+        meta_dir = tmp / "meta"
+        meta_dir.mkdir(exist_ok=True)
+        with open(meta_dir / "info.json", "w") as f:
+            json.dump(info, f, indent=2)
+
+        api.upload_folder(
+            folder_path=str(meta_dir), path_in_repo="meta",
+            repo_id=HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN,
+        )
+
+        result = {
+            "success": True,
+            "episode_index": episode_index,
+            "num_frames": num_frames,
+            "parquet_uploaded": True,
+            "video_uploaded": video_uploaded,
+            "dataset_url": f"https://huggingface.co/datasets/{HF_DATASET_REPO}",
+        }
+        return json.dumps(result)
+
+    finally:
+        shutil.rmtree(tmp, ignore_errors=True)
+
+
+def build_default_info(episode_index, num_frames):
+    return {
+        "codebase_version": "v2.0",
+        "robot_type": "unknown",
+        "total_episodes": episode_index + 1,
+        "total_frames": num_frames,
+        "total_tasks": 1,
+        "total_videos": 1,
+        "total_chunks": 1,
+        "chunks_size": CHUNKS_SIZE,
+        "fps": 30,
+        "splits": {"train": f"0:{episode_index + 1}"},
+        "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet",
+        "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4",
+        "features": {
+            "observation.camera_pose": {"dtype": "float32", "shape": [6],
+                                         "names": ["x", "y", "z", "roll", "pitch", "yaw"]},
+            "observation.left_hand": {"dtype": "float32", "shape": [9],
+                                       "names": ["wrist_x", "wrist_y", "wrist_z", "thumb_x", "thumb_y", "thumb_z",
+                                                  "index_x", "index_y", "index_z"]},
+            "observation.right_hand": {"dtype": "float32", "shape": [9],
+                                        "names": ["wrist_x", "wrist_y", "wrist_z", "index_x", "index_y", "index_z",
+                                                   "middle_x", "middle_y", "middle_z"]},
+            "action.camera_delta": {"dtype": "float32", "shape": [6],
+                                     "names": ["dx", "dy", "dz", "droll", "dpitch", "dyaw"]},
+            "action.left_hand_delta": {"dtype": "float32", "shape": [9],
+                                        "names": ["wrist_dx", "wrist_dy", "wrist_dz", "thumb_dx", "thumb_dy",
+                                                   "thumb_dz", "index_dx", "index_dy", "index_dz"]},
+            "action.right_hand_delta": {"dtype": "float32", "shape": [9],
+                                         "names": ["wrist_dx", "wrist_dy", "wrist_dz", "index_dx", "index_dy",
+                                                    "index_dz", "middle_dx", "middle_dy", "middle_dz"]},
+            "language_instruction": {"dtype": "string", "shape": [1], "names": None},
+            "timestamp": {"dtype": "float64", "shape": [1], "names": None},
+            "frame_index": {"dtype": "int64", "shape": [1], "names": None},
+            "episode_index": {"dtype": "int64", "shape": [1], "names": None},
+            "next.done": {"dtype": "bool", "shape": [1], "names": None},
+            "rgb": {"dtype": "video", "shape": [480, 640, 3],
+                    "names": ["height", "width", "channels"],
+                    "video_info": {"video.fps": 30, "video.codec": "h264",
+                                   "video.pix_fmt": "yuv420p", "video.is_depth_map": False,
+                                   "has_audio": False}},
+        },
+        "videos": {
+            "rgb": {"video_info": {"video.fps": 30, "video.codec": "h264",
+                                    "video.pix_fmt": "yuv420p", "video.is_depth_map": False,
+                                    "has_audio": False}}
+        },
+    }
+
+
+# Gradio UI (also exposes /api/convert endpoint automatically)
+demo = gr.Interface(
+    fn=convert_episode,
+    inputs=gr.Textbox(label="Episode JSON", lines=10, placeholder="Paste episode JSON here..."),
+    outputs=gr.Textbox(label="Result"),
+    title="DI LeRobot Converter",
+    description="Converts episode data from DI iOS app to LeRobot v2.0 format and uploads to HuggingFace dataset repo.",
+    api_name="convert",
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)