Untitled

#!/usr/bin/env python3
"""
Convert EgoDex test subset to a LeRobot dataset using only:
  - egocentric video
  - 48D hand actions

48D action layout per frame:
  left hand:  wrist pos (3) + wrist rot6d (6) + 5 fingertips xyz (15) = 24
  right hand: wrist pos (3) + wrist rot6d (6) + 5 fingertips xyz (15) = 24
  total = 48

Expected EgoDex layout:
  <egodex_root>/test/<task_name>/<episode_id>.mp4
  <egodex_root>/test/<task_name>/<episode_id>.hdf5

Example:
  python convert_egodex_test_to_lerobot.py \
      --egodex-root /data/egodex \
      --output-root /data/egodex_lerobot \
      --repo-id local/egodex_test_48d \
      --fps 30
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Iterator

import cv2
import h5py
import numpy as np


FINGERTIP_JOINTS = {
    "left": [
        "leftThumbTip",
        "leftIndexFingerTip",
        "leftMiddleFingerTip",
        "leftRingFingerTip",
        "leftLittleFingerTip",
    ],
    "right": [
        "rightThumbTip",
        "rightIndexFingerTip",
        "rightMiddleFingerTip",
        "rightRingFingerTip",
        "rightLittleFingerTip",
    ],
}


def rotation_matrix_to_6d(rot: np.ndarray) -> np.ndarray:
    """
    Convert a 3x3 rotation matrix to the 6D representation used in the H_RDT example:
    first two columns concatenated.
    """
    rot = np.asarray(rot, dtype=np.float32)
    return np.concatenate([rot[:, 0], rot[:, 1]], axis=0).astype(np.float32)


def construct_48d_action(transforms_group: h5py.Group, frame_idx: int) -> np.ndarray:
    """
    Build the 48D action vector for a single frame from EgoDex transforms.
    """
    action_vector: list[np.ndarray] = []

    for hand_side in ("left", "right"):
        hand_key = f"{hand_side}Hand"
        hand_transform = np.asarray(transforms_group[hand_key][frame_idx], dtype=np.float32)  # (4, 4)

        wrist_pos = hand_transform[:3, 3]
        wrist_rot6d = rotation_matrix_to_6d(hand_transform[:3, :3])

        action_vector.append(wrist_pos)
        action_vector.append(wrist_rot6d)

        for fingertip in FINGERTIP_JOINTS[hand_side]:
            fingertip_transform = np.asarray(transforms_group[fingertip][frame_idx], dtype=np.float32)
            fingertip_pos = fingertip_transform[:3, 3]
            action_vector.append(fingertip_pos)

    out = np.concatenate(action_vector, axis=0).astype(np.float32)
    if out.shape != (48,):
        raise ValueError(f"Expected action shape (48,), got {out.shape}")
    return out


def iter_video_frames_bgr(video_path: Path) -> Iterator[np.ndarray]:
    """
    Yield video frames as HxWx3 uint8 RGB arrays.
    """
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise RuntimeError(f"Failed to open video: {video_path}")

    try:
        while True:
            ok, frame_bgr = cap.read()
            if not ok:
                break
            frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
            yield frame_rgb
    finally:
        cap.release()


def get_video_info(video_path: Path) -> tuple[int, int, int, float]:
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        raise RuntimeError(f"Failed to open video: {video_path}")
    try:
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = float(cap.get(cv2.CAP_PROP_FPS))
        return frame_count, height, width, fps
    finally:
        cap.release()


def safe_hdf5_attr_str(f: h5py.File, key: str, default: str = "") -> str:
    if key not in f.attrs:
        return default
    value = f.attrs[key]
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    return str(value)


def get_instruction(f: h5py.File, task_name: str) -> str:
    """
    Pick the best available language string from EgoDex metadata.
    """
    llm_description = safe_hdf5_attr_str(f, "llm_description", "")
    llm_description2 = safe_hdf5_attr_str(f, "llm_description2", "")
    which = safe_hdf5_attr_str(f, "which_llm_description", "")

    if which == "2" and llm_description2:
        return llm_description2
    if llm_description:
        return llm_description
    if llm_description2:
        return llm_description2
    return task_name.replace("_", " ")


def collect_test_episodes(test_root: Path) -> list[tuple[str, Path, Path]]:
    """
    Returns a list of:
      (task_name, mp4_path, hdf5_path)
    """
    episodes: list[tuple[str, Path, Path]] = []

    if not test_root.exists():
        raise FileNotFoundError(f"Test root not found: {test_root}")

    for task_dir in sorted(p for p in test_root.iterdir() if p.is_dir()):
        task_name = task_dir.name
        mp4s = sorted(task_dir.glob("*.mp4"))

        for mp4_path in mp4s:
            hdf5_path = mp4_path.with_suffix(".hdf5")
            if not hdf5_path.exists():
                print(f"[WARN] Missing HDF5 for {mp4_path}, skipping")
                continue
            episodes.append((task_name, mp4_path, hdf5_path))

    return episodes


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--egodex-root", type=Path, required=True, help="Root containing the EgoDex 'test' directory")
    parser.add_argument("--output-root", type=Path, required=True, help="Local output root for LeRobot dataset")
    parser.add_argument("--repo-id", type=str, required=True, help="LeRobot dataset repo_id, e.g. local/egodex_test_48d")
    parser.add_argument("--fps", type=float, default=30.0, help="Target FPS in LeRobot metadata")
    parser.add_argument("--robot-type", type=str, default="egodex_hands", help="Robot type string for LeRobot metadata")
    parser.add_argument("--video-key", type=str, default="observation.images.ego", help="Feature key for ego video")
    parser.add_argument("--action-key", type=str, default="action", help="Feature key for 48D actions")
    parser.add_argument("--task-as-frame-field", action="store_true",
                        help="Also include frame['task']=... for compatibility with some LeRobot installs")
    args = parser.parse_args()

    test_root = args.egodex_root / "test"
    episodes = collect_test_episodes(test_root)

    if not episodes:
        raise RuntimeError(f"No test episodes found under {test_root}")

    args.output_root.mkdir(parents=True, exist_ok=True)

    # Import here so the script can still be inspected without LeRobot installed.
    from lerobot.datasets.lerobot_dataset import LeRobotDataset

    # Use the first video to define image shape.
    _, first_mp4, first_h5 = episodes[0]
    _, height, width, detected_fps = get_video_info(first_mp4)

    print(f"Found {len(episodes)} episodes")
    print(f"Video shape: {height}x{width}x3")
    print(f"First-video FPS: {detected_fps:.3f}")
    print(f"LeRobot metadata FPS: {args.fps}")

    features = {
        args.video_key: {
            "dtype": "video",
            "shape": (height, width, 3),
            "names": ["height", "width", "channel"],
        },
        args.action_key: {
            "dtype": "float32",
            "shape": (48,),
            "names": [f"a_{i}" for i in range(48)],
        },
        # Optional but useful for alignment/debugging.
        "timestamp": {
            "dtype": "float32",
            "shape": (1,),
            "names": None,
        },
    }

    dataset = LeRobotDataset.create(
        repo_id=args.repo_id,
        root=args.output_root,
        fps=args.fps,
        robot_type=args.robot_type,
        features=features,
        use_videos=True,
    )

    summary = {
        "repo_id": args.repo_id,
        "num_episodes": 0,
        "episodes": [],
    }

    for episode_idx, (task_name, mp4_path, hdf5_path) in enumerate(episodes):
        print(f"[{episode_idx + 1}/{len(episodes)}] {task_name} / {mp4_path.name}")

        with h5py.File(hdf5_path, "r") as f:
            if "transforms" not in f:
                print(f"[WARN] Missing 'transforms' in {hdf5_path}, skipping")
                continue

            transforms_group = f["transforms"]
            instruction = get_instruction(f, task_name)

            # Determine HDF5 length from one transform dataset.
            first_key = next(iter(transforms_group.keys()))
            hdf5_num_frames = int(transforms_group[first_key].shape[0])

            video_frames = list(iter_video_frames_bgr(mp4_path))
            video_num_frames = len(video_frames)

            if video_num_frames == 0:
                print(f"[WARN] Empty video: {mp4_path}, skipping")
                continue

            num_frames = min(video_num_frames, hdf5_num_frames)
            if num_frames != video_num_frames or num_frames != hdf5_num_frames:
                print(
                    f"[WARN] Frame mismatch for {mp4_path.name}: "
                    f"video={video_num_frames}, hdf5={hdf5_num_frames}, using {num_frames}"
                )

            for frame_idx in range(num_frames):
                image = video_frames[frame_idx]
                action = construct_48d_action(transforms_group, frame_idx)
                timestamp = np.array([frame_idx / float(args.fps)], dtype=np.float32)

                frame = {
                    args.video_key: image,
                    args.action_key: action,
                    "timestamp": timestamp,
                }

                if args.task_as_frame_field:
                    frame["task"] = instruction

                # Some LeRobot versions accept task as a kwarg, which is the cleaner route.
                try:
                    dataset.add_frame(frame, task=instruction)
                except TypeError:
                    dataset.add_frame(frame)

            dataset.save_episode()

            summary["num_episodes"] += 1
            summary["episodes"].append(
                {
                    "episode_index": episode_idx,
                    "task_name": task_name,
                    "video": str(mp4_path),
                    "hdf5": str(hdf5_path),
                    "frames_used": num_frames,
                    "instruction": instruction,
                }
            )

    dataset.finalize()

    with open(args.output_root / "conversion_summary.json", "w", encoding="utf-8") as fp:
        json.dump(summary, fp, indent=2, ensure_ascii=False)

    print("\nDone.")
    print(f"LeRobot dataset written to: {args.output_root}")
    print(f"Episodes saved: {summary['num_episodes']}")


if __name__ == "__main__":
    main()