scripts/data/prepare_libero_gear.py

#!/usr/bin/env python3
"""
Prepare LIBERO dataset for DreamZero GEAR pipeline.

Converts chunked LeRobot v2 format (multiple episodes per parquet, images as PNG bytes)
into individual-episode format (one parquet + one mp4 per episode) expected by
convert_lerobot_to_gear.py.

Usage:
    python3 prepare_libero_gear.py \
        --input-dir /root/autodl-tmp/data/libero \
        --output-dir /root/autodl-tmp/data/libero_gear
"""

import argparse
import json
import os
import io
import sys
import time
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from PIL import Image
import cv2


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare LIBERO data for DreamZero GEAR pipeline")
    parser.add_argument("--input-dir", required=True, help="Path to original LIBERO dataset")
    parser.add_argument("--output-dir", required=True, help="Path for output GEAR-ready dataset")
    parser.add_argument("--num-workers", type=int, default=4, help="Number of parallel workers")
    parser.add_argument("--skip-video", action="store_true", help="Skip video encoding (test only)")
    return parser.parse_args()


def load_info(info_path: Path) -> dict:
    with open(info_path) as f:
        return json.load(f)


def save_info(info: dict, output_path: Path, num_episodes: int, total_frames: int):
    """Update info.json for individual-episode format."""
    info["total_episodes"] = num_episodes
    info["total_frames"] = total_frames
    info["chunks_size"] = 2000  # All episodes in chunk-000
    info["data_path"] = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet"
    info["video_path"] = "videos/{video_key}/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.mp4"

    # Remove meta/episodes path reference since we'll embed tasks directly
    info.pop("splits", None)

    with open(output_path / "meta" / "info.json", "w") as f:
        json.dump(info, f, indent=2)


def decode_png_to_rgb(png_bytes: bytes) -> np.ndarray:
    """Decode PNG bytes to RGB numpy array (H, W, 3) uint8."""
    img = Image.open(io.BytesIO(png_bytes))
    return np.array(img.convert("RGB"))


def extract_episode_metadata(input_dir: Path) -> tuple[pd.DataFrame, dict]:
    """Read tasks from the episodes metadata."""
    meta_dir = input_dir / "meta" / "episodes"
    if meta_dir.exists():
        ep_files = sorted(meta_dir.rglob("*.parquet"))
        if ep_files:
            df = pd.read_parquet(ep_files[0])
            tasks = {}
            for _, row in df.iterrows():
                ep_idx = row["episode_index"]
                tasks[ep_idx] = row["tasks"]
            return df, tasks

    # Fallback: scan parquet files for task_index
    return None, {}


def process_parquet_file(
    parquet_path: Path,
    output_data_dir: Path,
    output_video_dir: Path,
    fps: float,
    skip_video: bool = False,
) -> tuple[int, int]:
    """
    Process a single chunked parquet file.
    Returns (num_episodes_processed, num_frames_processed).
    """
    # Read the parquet file
    df = pd.read_parquet(parquet_path)

    # Group by episode_index
    episodes_processed = 0
    frames_processed = 0

    for ep_idx, group in df.groupby("episode_index"):
        ep_idx = int(ep_idx)
        group = group.reset_index(drop=True)
        n_frames = len(group)

        # Output parquet path
        ep_parquet_path = output_data_dir / f"episode_{ep_idx:06d}.parquet"

        # Drop the image columns for the parquet (they're in the video now)
        # But keep them for now — the official stats computation only uses numeric columns
        # We need to keep image columns as they might be needed by the dataset loader
        # Actually, for the GEAR format, images should ONLY be in videos.
        # Remove image columns to avoid confusion.
        parquet_cols = [c for c in group.columns
                        if not c.startswith("observation.images.")]
        df_out = group[parquet_cols].copy()

        # Write parquet
        table = pa.Table.from_pandas(df_out)
        pq.write_table(table, ep_parquet_path)

        if not skip_video:
            # Decode and write video for observation.images.image (first camera)
            frames = []
            for _, row in group.iterrows():
                img_bytes = row["observation.images.image"]["bytes"]
                frame = decode_png_to_rgb(img_bytes)
                frames.append(frame)

            # Write mp4 video
            ep_video_path = output_video_dir / f"episode_{ep_idx:06d}.mp4"
            height, width = frames[0].shape[:2]
            fourcc = cv2.VideoWriter_fourcc(*"mp4v")
            out = cv2.VideoWriter(
                str(ep_video_path), fourcc, fps, (width, height)
            )
            for frame in frames:
                # cv2 uses BGR order
                out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
            out.release()

        episodes_processed += 1
        frames_processed += n_frames

    return episodes_processed, frames_processed


def main():
    args = parse_args()
    input_dir = Path(args.input_dir)
    output_dir = Path(args.output_dir)

    # Create output structure
    output_data_dir = output_dir / "data" / "chunk-000"
    output_video_dir = output_dir / "videos" / "observation.images.image" / "chunk-000"
    output_meta_dir = output_dir / "meta"

    output_data_dir.mkdir(parents=True, exist_ok=True)
    output_video_dir.mkdir(parents=True, exist_ok=True)
    output_meta_dir.mkdir(parents=True, exist_ok=True)

    # Load original info.json
    info = load_info(input_dir / "meta" / "info.json")
    fps = info.get("fps", 10.0)

    # Find all data parquet files
    data_dir = input_dir / "data" / "chunk-000"
    parquet_files = sorted(data_dir.glob("file-*.parquet"))
    print(f"Found {len(parquet_files)} parquet files")

    # Process each file
    total_episodes = 0
    total_frames = 0
    start_time = time.time()

    for i, pf in enumerate(parquet_files):
        n_eps, n_frames = process_parquet_file(
            pf, output_data_dir, output_video_dir, fps,
            skip_video=args.skip_video,
        )
        total_episodes += n_eps
        total_frames += n_frames

        elapsed = time.time() - start_time
        rate = (i + 1) / elapsed if elapsed > 0 else 0
        eta = (len(parquet_files) - i - 1) / rate if rate > 0 else 0
        print(
            f"  [{i+1}/{len(parquet_files)}] {pf.name}: "
            f"{n_eps} eps, {n_frames} frames "
            f"({rate:.1f} files/min, ETA {eta/60:.0f}min)"
        )

    # Write info.json
    save_info(info, output_dir, total_episodes, total_frames)

    # Copy tasks metadata if available
    meta_ep_dir = input_dir / "meta" / "episodes" / "chunk-000"
    if meta_ep_dir.exists():
        ep_files = sorted(meta_ep_dir.glob("*.parquet"))
        if ep_files:
            ep_meta_df = pd.read_parquet(ep_files[0])
            # Extract task_index → task mapping
            tasks = {}
            for _, row in ep_meta_df.iterrows():
                ep_idx = int(row["episode_index"])
                task_text = row["tasks"]
                if isinstance(task_text, np.ndarray):
                    task_text = task_text.item() if task_text.size > 0 else ""
                elif isinstance(task_text, bytes):
                    task_text = task_text.decode("utf-8")
                tasks[ep_idx] = str(task_text)

            # Write tasks.jsonl for GEAR format
            unique_tasks = sorted(set(tasks.values()))
            with open(output_meta_dir / "tasks.jsonl", "w") as f:
                for ti, task in enumerate(unique_tasks):
                    f.write(json.dumps({"task_index": ti, "task": task}) + "\n")

            # Write episodes.jsonl
            with open(output_meta_dir / "episodes.jsonl", "w") as f:
                for _, row in ep_meta_df.iterrows():
                    ep_idx = int(row["episode_index"])
                    length = int(row["length"])
                    task_text = tasks.get(ep_idx, "")
                    task_index = unique_tasks.index(task_text) if task_text in unique_tasks else -1
                    f.write(json.dumps({
                        "episode_index": ep_idx,
                        "length": length,
                        "task_index": task_index,
                    }) + "\n")

            print(f"Wrote {len(unique_tasks)} tasks and {total_episodes} episode entries")

    print(f"\nDone! {total_episodes} episodes, {total_frames} frames")
    print(f"Output: {output_dir}")
    print(f"Time: {(time.time() - start_time)/60:.1f} minutes")
    print(f"\nNext step: run convert_lerobot_to_gear.py on the output dir")


if __name__ == "__main__":
    main()