#!/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()