lerobot/examples/behavior_1k/convert_to_lerobot_v3.py

#!/usr/bin/env python

"""
Convert a single BEHAVIOR-1K task from HDF5 to LeRobotDataset v3.0 format.

Usage examples:
# Convert a single task
python convert_to_lerobot_v3.py \
    --data-folder /path/to/data \
    --repo-id "username/behavior-1k-assembling-gift-baskets" \
    --task-id 0 \
    --push-to-hub

"""

import argparse
import logging
import os
from pathlib import Path

import h5py
import numpy as np
from tqdm import tqdm

from lerobot.utils.utils import init_logging

from .behavior_lerobot_dataset_v3 import BehaviorLeRobotDatasetV3
from .behaviour_1k_constants import BEHAVIOR_DATASET_FEATURES, FPS, ROBOT_TYPE, TASK_INDICES_TO_NAMES

init_logging()


def load_hdf5_episode(hdf5_path: str, episode_id: int = 0) -> dict:
    """
    Load episode data from HDF5 file.

    Args:
        hdf5_path: Path to the HDF5 file
        episode_id: Episode ID to load (default: 0)

    Returns:
        Dictionary containing episode data
    """
    episode_data = {}

    with h5py.File(hdf5_path, "r") as f:
        # Find the episode with most samples if episode_id not specified
        if episode_id == -1:
            num_samples = [f["data"][key].attrs["num_samples"] for key in f["data"]]
            episode_id = num_samples.index(max(num_samples))

        demo_key = f"demo_{episode_id}"
        if demo_key not in f["data"]:
            raise ValueError(f"Episode {episode_id} not found in {hdf5_path}")

        demo_data = f["data"][demo_key]

        # Load actions
        episode_data["action"] = np.array(demo_data["action"][:])

        # Load observations
        episode_data["obs"] = {}
        for key in demo_data["obs"]:
            episode_data["obs"][key] = np.array(demo_data["obs"][key][:])

        # Load attributes
        episode_data["attrs"] = {}
        for attr_name in demo_data.attrs:
            episode_data["attrs"][attr_name] = demo_data.attrs[attr_name]

        # Add global attributes
        for attr_name in f["data"].attrs:
            episode_data["attrs"][f"global_{attr_name}"] = f["data"].attrs[attr_name]

    return episode_data


def convert_episode(
    data_folder: str,
    task_id: int,
    demo_id: int,
    dataset: BehaviorLeRobotDatasetV3,
    include_videos: bool = True,
    include_segmentation: bool = True,
) -> None:
    """
    Convert a single episode from HDF5 to LeRobotDataset v3.0 format.

    Args:
        data_folder: Base data folder containing HDF5 files
        repo_id: Repository ID for the dataset
        task_id: Task ID
        demo_id: Demo ID (episode ID)
        dataset: BehaviorLeRobotDatasetV3 instance to add data to
        include_videos: Whether to include video data
        include_segmentation: Whether to include segmentation data
    """
    # Construct paths
    task_name = TASK_INDICES_TO_NAMES[task_id]
    hdf5_path = f"{data_folder}/2025-challenge-rawdata/task-{task_id:04d}/episode_{demo_id:08d}.hdf5"

    if not os.path.exists(hdf5_path):
        logging.error(f"HDF5 file not found: {hdf5_path}")
        return

    logging.info(f"Converting episode {demo_id} from task {task_name}")

    # Load episode data
    episode_data = load_hdf5_episode(hdf5_path, episode_id=0)

    # Filter out segmentation if not requested
    if not include_segmentation:
        keys_to_remove = [k for k in episode_data["obs"] if "seg_instance_id" in k]
        for key in keys_to_remove:
            del episode_data["obs"][key]

    # Add episode to dataset
    dataset.add_episode_from_hdf5(
        hdf5_data=episode_data,
        task_id=task_id,
        episode_id=demo_id,
        include_videos=include_videos,
    )


def convert_task_to_dataset(
    data_folder: str,
    repo_id: str,
    task_id: int,
    push_to_hub: bool = False,
) -> None:
    """
    Convert a single BEHAVIOR-1K task from HDF5 to LeRobotDataset v3.0 format.

    Args:
        data_folder: Base folder containing HDF5 data
        repo_id: Repository ID (e.g., "username/behavior-1k-task-name")
        task_id: Task ID to convert
        push_to_hub: Whether to push to HuggingFace Hub
    """
    task_name = TASK_INDICES_TO_NAMES[task_id]
    task_folder = f"{data_folder}/2025-challenge-rawdata/task-{task_id:04d}"

    if not os.path.exists(task_folder):
        raise ValueError(f"Task folder not found: {task_folder}")

    # Create output directory
    output_dir = Path.home() / ".cache/huggingface/lerobot" / repo_id
    output_dir.mkdir(parents=True, exist_ok=True)

    logging.info(f"Converting task '{task_name}' (ID: {task_id}) to: {output_dir}")

    # Initialize dataset for this task
    dataset = BehaviorLeRobotDatasetV3.create(
        repo_id=repo_id,
        fps=FPS,
        features=BEHAVIOR_DATASET_FEATURES,
        robot_type=ROBOT_TYPE,
    )

    # Find all episodes in the task folder
    task_episode_ids = []
    for filename in os.listdir(task_folder):
        if filename.startswith("episode_") and filename.endswith(".hdf5"):
            eid = int(filename.split("_")[1].split(".")[0])
            task_episode_ids.append(eid)
    task_episode_ids.sort()

    logging.info(f"Processing {len(task_episode_ids)} episodes for task {task_name}")

    # Convert each episode
    episodes_converted = 0
    for demo_id in tqdm(task_episode_ids, desc="Converting episodes"):
        convert_episode(
            data_folder=data_folder,
            task_id=task_id,
            demo_id=demo_id,
            dataset=dataset,
            include_videos=True,
            include_segmentation=True,
        )
        episodes_converted += 1

    logging.info(f"Converted {episodes_converted} episodes for task {task_name}")

    # Finalize dataset
    logging.info(f"Finalizing dataset for task {task_name}...")
    dataset.finalize()

    # Push to hub if requested
    if push_to_hub:
        logging.info(f"Pushing task {task_name} dataset to HuggingFace Hub...")
        dataset.push_to_hub()

    logging.info("Conversion complete!")


def main():
    parser = argparse.ArgumentParser(description="Convert a single BEHAVIOR-1K task to LeRobotDataset v3.0")
    parser.add_argument("--data-folder", type=str, required=True, help="Path to the data folder")
    parser.add_argument(
        "--repo-id",
        type=str,
        required=True,
        help="Output repository ID (e.g., 'username/behavior-1k-assembling-gift-baskets')",
    )
    parser.add_argument(
        "--task-id", type=int, required=True, help="Task ID to convert (e.g., 0 for assembling_gift_baskets)"
    )
    parser.add_argument(
        "--push-to-hub", action="store_true", help="Push dataset to HuggingFace Hub after conversion"
    )

    args = parser.parse_args()

    # Convert single task to dataset
    convert_task_to_dataset(
        data_folder=args.data_folder,
        repo_id=args.repo_id,
        task_id=args.task_id,
        push_to_hub=args.push_to_hub,
    )


if __name__ == "__main__":
    main()