""" This script converts the LeRobot Hugging Face dataset format to the CV Processed format. It extracts frames from videos, processes low-dimensional data, and organizes the output in the CV Processed structure. Input (LeRobot dataset format): - data/chunk-*/episode_{episode_num}.parquet - videos/chunk-*/camera/{episode_num}.mp4 - meta/episodes.jsonl - meta/info.json Output (CV Processed format): - task_name/episode_{episode_num}/ - rgb/camera_{camera_num}/frame_{frame_num}.jpg - depth/camera_{camera_num}/frame_{frame_num}.npz - lowdim/camera_{camera_num}/frame_{frame_num}.npz AGIBot dataset path: s3://robotics-manip-lbm/agibot_world/agibot_world_alpha/lerobot Example usage: python3 convert_lerobot_to_cvprocessed.py --dataset_path lerobot/datasets/agibot --output_path cv_processed/agibot --num_workers 8 --tmp_dir /tmp """ import argparse import io import json import os import re import subprocess import tempfile import traceback from concurrent.futures import ThreadPoolExecutor from pathlib import Path import numpy as np import pyarrow.parquet as pq import logging from datetime import datetime from typing import Dict, List, Tuple from PIL import Image from file_utils import copy_to_temp_file, file_exists, json_load, list_directory # Image dimensions DEFAULT_IMG_H, DEFAULT_IMG_W = 480, 640 def setup_logging(log_dir="logs"): os.makedirs(log_dir, exist_ok=True) log_name = datetime.now().strftime("lerobot_to_cvprocessed_%Y%m%d_%H%M%S.log") log_path = os.path.join(log_dir, log_name) logger = logging.getLogger() logger.setLevel(logging.INFO) formatter = logging.Formatter( "[%(asctime)s][%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S" ) # Stream handler for terminal sh = logging.StreamHandler() sh.setFormatter(formatter) logger.addHandler(sh) # File handler for log file fh = logging.FileHandler(log_path) fh.setFormatter(formatter) logger.addHandler(fh) logging.info(f"Logging to {log_path}") def parse_args(): parser = argparse.ArgumentParser() # Input/Output paths parser.add_argument( "--dataset_path", type=str, required=True, help="All other paths become relative to this.", ) parser.add_argument( "--output_path", type=str, required=True, help="Path to save the processed dataset.", ) parser.add_argument( "--tmp_dir", type=str, default=None, help="Directory for tempfiles. If None, defaults to /tmp", ) # Processing parameters parser.add_argument( "--fps", type=float, default=None, help="Auto-detected from info_path if not specified", ) parser.add_argument( "--num_workers", type=int, default=None, help="Number of worker threads (default: auto-detect)", ) # File/column naming (probably no need to change any of these for standard LeRobot) parser.add_argument("--meta_episodes_path", type=str, default="meta/episodes.jsonl") parser.add_argument("--info_path", type=str, default="meta/info.json") parser.add_argument("--data_path", type=str, default="data") parser.add_argument("--videos_path", type=str, default="videos") # Lowdim columns can vary across LeRobot datasets so these need to be customized parser.add_argument( "--lowdim_columns", type=str, nargs="+", default=[ "task_index", "action", "observation.state", "intrinsics", "extrinsics_aligned", "timestamp", ], ) parser.add_argument("--frame_index_col", type=str, default="frame_index") parser.add_argument("--episode_index_col", type=str, default="episode_index") parser.add_argument( "--depth_columns", type=str, nargs="+", default=["observation.images.cam_top_depth"], ) parser.add_argument( "--episode_file_pattern", type=str, default="episode_{:06d}.parquet" ) parser.add_argument("--video_file_pattern", type=str, default="episode_{:06d}.mp4") parser.add_argument( "--debug", action="store_true", help="If set, process only one episode sequentially for debugging.", ) return parser.parse_args() def resolve_path(base_path: str, relative_path: str) -> str: # If relative_path is already absolute (starts with s3:// or /), return as-is if relative_path.startswith("s3://") or relative_path.startswith("/"): return relative_path return f"{base_path.rstrip('/')}/{relative_path.lstrip('/')}" def detect_fps(info_path: str) -> float: """Automatically detect FPS from info.json file.""" info = json_load(info_path) if "fps" in info: fps = float(info["fps"]) logging.info(f"Detected FPS from info.json: {fps}") else: fps = 30.0 logging.info( f"Warning: Could not detect FPS from info.json, using default FPS {fps}" ) return fps def discover_chunks(base_path: str, pattern: str = "chunk-*") -> List[str]: """Discover all chunk directories in the base path.""" chunk_dirs = [] for item in list_directory(base_path): if item.startswith("chunk-"): chunk_path = f"{base_path.rstrip('/')}/{item}" chunk_dirs.append(chunk_path) chunk_dirs.sort() # Sort to ensure consistent ordering logging.info( f"Discovered {len(chunk_dirs)} chunks in {base_path}: {[os.path.basename(d) for d in chunk_dirs]}" ) return chunk_dirs def discover_cameras(video_chunks: List[str]) -> Dict[str, str]: """Discover available cameras by scanning video chunk directories.""" cameras = {} if not video_chunks: return cameras first_chunk = video_chunks[0] for item in list_directory(first_chunk): camera_name = item.split(".")[-1] if "." in item else item cameras[camera_name] = item logging.info(f"Discovered cameras: {list(cameras.keys())}") return cameras def build_episode_lookup_chunk(chunk_dir: str) -> Dict[int, str]: """Build episode lookup for a single chunk directory.""" episode_lookup = {} for file in list_directory(chunk_dir): if file.endswith(".parquet") and "episode_" in file: match = re.search(r"episode_(\d+)\.parquet", file) if match: ep_num = int(match.group(1)) episode_lookup[ep_num] = f"{chunk_dir.rstrip('/')}/{file}" return episode_lookup def build_episode_lookup(data_chunks: List[str]) -> Dict[int, str]: logging.info(f"Building episode lookup dict from {len(data_chunks)} chunks...") # Process chunks using ThreadPoolExecutor episode_lookup = {} with ThreadPoolExecutor() as executor: chunk_results = list(executor.map(build_episode_lookup_chunk, data_chunks)) # Merge results for chunk_result in chunk_results: episode_lookup.update(chunk_result) return episode_lookup def build_video_lookup_chunk( chunk_dir_cameras: Tuple[str, Dict[str, str]], ) -> Dict[Tuple[int, str], str]: """Build video lookup for a single chunk directory.""" chunk_dir, cameras = chunk_dir_cameras video_lookup = {} for _camera_name, camera_path in cameras.items(): camera_dir = f"{chunk_dir.rstrip('/')}/{camera_path}" files = list_directory(camera_dir) for file in files: if file.endswith(".mp4") and "episode_" in file: match = re.search(r"episode_(\d+)\.mp4", file) if match: ep_num = int(match.group(1)) video_lookup[(ep_num, camera_path)] = f"{camera_dir}/{file}" return video_lookup def build_video_lookup( video_chunks: List[str], cameras: Dict[str, str] ) -> Dict[Tuple[int, str], str]: logging.info( f"Building video lookup dict from {len(video_chunks)} chunks and {len(cameras)} cameras..." ) # Process chunks using ThreadPoolExecutor video_lookup = {} chunk_camera_pairs = [(chunk, cameras) for chunk in video_chunks] with ThreadPoolExecutor() as executor: chunk_results = list(executor.map(build_video_lookup_chunk, chunk_camera_pairs)) # Merge results for chunk_result in chunk_results: video_lookup.update(chunk_result) return video_lookup def discover_metadata(meta_episodes_path: str) -> Dict[int, dict]: """Discover metadata from the episodes.jsonl file.""" metadata = {} if not file_exists(meta_episodes_path): logging.warning(f"Metadata file {meta_episodes_path} does not exist.") return metadata with open(meta_episodes_path, "r") as f: for line in f: entry = json.loads(line) episode_index = entry.get("episode_index") if episode_index is not None: metadata[episode_index] = entry return metadata def extract_single_frame_ffmpeg( video_path: str, frame_index: int, output_path: str, fps: float = 30.0 ) -> bool: """Extract a specific frame from video using ffmpeg. Video_path should be local.""" try: timestamp = frame_index / fps cmd = [ "ffmpeg", "-ss", str(timestamp), "-i", video_path, "-vframes", "1", "-y", output_path, "-loglevel", "quiet", ] with open(os.devnull, "w") as devnull: result = subprocess.run( cmd, stdout=devnull, stderr=devnull, stdin=subprocess.DEVNULL, timeout=10, ) return result.returncode == 0 and file_exists(output_path) except ( subprocess.TimeoutExpired, FileNotFoundError, subprocess.CalledProcessError, ): return False def extract_multiple_frames_ffmpeg( video_path: str, frame_requests: List[Tuple[int, str]], fps: float = 30.0 ) -> Dict[int, bool]: """Extract multiple frames from a video efficiently using ffmpeg filter.""" if not frame_requests: return {} results = {} try: # Sort requests by frame index for efficiency sorted_requests = sorted(frame_requests, key=lambda x: x[0]) # Use ffmpeg select filter to extract frames in batches of 10 to avoid command line length limits batch_size = 10 for i in range(0, len(sorted_requests), batch_size): batch = sorted_requests[i : i + batch_size] # Build filter for selecting multiple frames frame_numbers = [str(req[0]) for req in batch] select_filter = ( f"select='{'+'.join([f'eq(n,{num})' for num in frame_numbers])}'" ) cmd = [ "ffmpeg", "-i", video_path, "-vf", select_filter, "-vsync", "0", # Don't duplicate frames "-y", str(Path(batch[0][1]).parent / f"batch_{i}_%d.jpg"), "-loglevel", "quiet", ] with open(os.devnull, "w") as devnull: result = subprocess.run( cmd, stdout=devnull, stderr=devnull, stdin=subprocess.DEVNULL, timeout=30, ) # Rename extracted files to target names and check success for j, (frame_index, target_path) in enumerate(batch): extracted_path = Path(batch[0][1]).parent / f"batch_{i}_{j + 1}.jpg" if result.returncode == 0 and extracted_path.exists(): try: os.rename(str(extracted_path), target_path) results[frame_index] = True except Exception: results[frame_index] = False if extracted_path.exists(): extracted_path.unlink() else: results[frame_index] = False except ( subprocess.TimeoutExpired, FileNotFoundError, subprocess.CalledProcessError, ) as e: logging.error( f"Batch ffmpeg extraction failed: {e}. Falling back to single frame extractions." ) # Fall back to individual extractions for frame_index, output_path in frame_requests: if frame_index not in results: results[frame_index] = extract_single_frame_ffmpeg( video_path, frame_index, output_path, fps ) return results def extract_intrinsics_matrix(intrinsics: np.ndarray) -> np.ndarray: """Convert intrinsics list to 3x3 matrix.""" # check if already 3x3 if intrinsics.shape == (3, 3): return intrinsics if intrinsics.size != 9: raise ValueError(f"Intrinsics list must have 9 elements, got {intrinsics.size}") # Raise warning if fx, fy are not reasonably close to each other fx, fy = intrinsics[0], intrinsics[4] if not np.isclose(fx, fy, rtol=5): logging.warning(f"Focal lengths are not close: fx={fx}, fy={fy}") return intrinsics.reshape(3, 3) def extract_extrinsics_matrix(rotation, translation) -> np.ndarray: """Convert extrinsics list to 4x4 matrix.""" extrinsics = np.eye(4) extrinsics[:3, :3] = rotation.reshape(3, 3) extrinsics[:3, 3] = translation if not np.isclose(np.linalg.det(extrinsics[:3, :3]), 1.0, rtol=1e-3): logging.warning( f"Determinant of rotation part is not close to 1: det={np.linalg.det(extrinsics[:3, :3])}" ) return extrinsics def detect_image_dimensions( info_path: str, cameras: dict ) -> dict: # return hxw for each camera """Automatically detect image dimensions from info.json file.""" info = json_load(info_path) camera_name_to_img_shape_lookup = {} for camera_name, camera_info in cameras.items(): if "shape" in info["features"][camera_info]: img_h = int(info["features"][camera_info]["shape"][0]) img_w = int(info["features"][camera_info]["shape"][1]) logging.info( f"Detected image dimensions for {camera_name} from info.json: H={img_h}, W={img_w}" ) else: img_h, img_w = DEFAULT_IMG_H, DEFAULT_IMG_W logging.info( f"Warning: Could not detect image dimensions from info.json, using default H={img_h}, W={img_w}" ) camera_name_to_img_shape_lookup[camera_name] = (img_h, img_w) return camera_name_to_img_shape_lookup def process_episode(episode_data: tuple): """Process a single episode and convert it to CV Processed format.""" ( ep_idx, episode_path, task_name, video_lookup, cameras, output_path, fps, metadata, camera_name_to_img_shape, config, ) = episode_data try: logging.info(f"Processing episode {ep_idx} at {episode_path}") # Read parquet file (download from S3 if needed) if episode_path.startswith("s3"): with copy_to_temp_file(episode_path) as temp_parquet: df = pq.read_table(temp_parquet).to_pandas() else: df = pq.read_table(episode_path).to_pandas() # Create output directories for this episode episode_output_dir = os.path.join( output_path, f"{task_name}", f"episode_{ep_idx:06d}" ) for i, camera in enumerate(cameras.keys()): rgb_output_dir = os.path.join(episode_output_dir, "rgb", f"camera_{i}") depth_output_dir = os.path.join(episode_output_dir, "depth", f"camera_{i}") lowdim_output_dir = os.path.join( episode_output_dir, "lowdim", f"camera_{i}" ) os.makedirs(rgb_output_dir, exist_ok=True) os.makedirs(depth_output_dir, exist_ok=True) os.makedirs(lowdim_output_dir, exist_ok=True) # Process video frames for RGB images primary_camera_name = camera # 'back_right_fisheye' primary_camera_path = cameras[ primary_camera_name ] # 'observation.images.back_right_fisheye' video_path = video_lookup.get( (ep_idx, primary_camera_path) ) if not video_path: logging.error( f"Could not find video for episode {ep_idx}, camera {primary_camera_name}" ) return False # Create a temporary directory for video processing with tempfile.TemporaryDirectory(dir=config["tmp_dir"]) as tmpdir: # Download video if it's on S3 if video_path.startswith("s3"): local_video_path = os.path.join( tmpdir, f"episode_{ep_idx}_video.mp4" ) cmd = f"aws s3 cp {video_path} {local_video_path}" result = subprocess.run( cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE ) if result.returncode != 0: logging.error( f"Failed to download video {video_path}: {result.stderr.decode().strip()}" ) return False else: local_video_path = video_path # Prepare frame extraction requests frame_requests = [] for idx, row in df.iterrows(): frame_idx = row[config["frame_index_col"]] output_file = os.path.join(rgb_output_dir, f"frame_{idx:06d}.jpg") frame_requests.append((frame_idx, output_file)) # Extract frames from video and check for failures results = extract_multiple_frames_ffmpeg( local_video_path, frame_requests, fps ) if not all(results.values()): failed_frames = [ idx for idx, success in results.items() if not success ] logging.error( f"Failed to extract some frames {failed_frames} from video {video_path}" ) # Save low-dimensional data for idx, row in df.iterrows(): lowdim_data = {} # Extract lowdim columns for col in config["lowdim_columns"]: if col in row: value = row[col] # Handle nested value fields with dot notation if isinstance(value, dict) and "." in col: parts = col.split(".") parent = parts[0] child = parts[1] if parent not in lowdim_data: lowdim_data[parent] = {} lowdim_data[parent][child] = value else: lowdim_data[col] = value elif col in ["intrinsics"]: camera_param_key = f"{col}.{primary_camera_name}" if camera_param_key in row: lowdim_data[col] = extract_intrinsics_matrix( row[camera_param_key] ) elif col in ["extrinsics_aligned"]: camera_param_rot_key = ( f"{col}.{primary_camera_name}.rotation_matrix" ) camera_param_trans_key = ( f"{col}.{primary_camera_name}.translation_vector" ) if ( camera_param_rot_key in row and camera_param_trans_key in row ): lowdim_data["pose"] = extract_extrinsics_matrix( row[camera_param_rot_key], row[camera_param_trans_key], ) else: logging.warning( f"Column {col} not found in episode {ep_idx}, frame {idx}" ) # TODO: Add shape, prompt data if needed lowdim_data["shape"] = camera_name_to_img_shape.get( primary_camera_name, (DEFAULT_IMG_H, DEFAULT_IMG_W) ) lowdim_data["prompt"] = metadata.get(ep_idx, {}).get("tasks", "") # import ipdb; ipdb.set_trace() # Save lowdim data as npz file output_file = os.path.join( lowdim_output_dir, f"frame_{idx:06d}.npz" ) np.savez(output_file, **lowdim_data) # lowdim_data.keys() # dict_keys(['task_index', 'action', 'observation.state', 'intrinsics', 'pose', 'shape', 'prompt']) # load and save depth for col in config["depth_columns"]: if col in row: depth_bytes = row[col]["bytes"] depth_pil_image = Image.open(io.BytesIO(depth_bytes)) depth_array = ( np.array(depth_pil_image).astype(np.float32) / 1000.0 ) # Convert mm to meters depth_output_file = os.path.join( depth_output_dir, f"frame_{idx:06d}.npz" ) np.savez_compressed(depth_output_file, data=depth_array) else: logging.warning( f"Depth column {col} not found in episode {ep_idx}, frame {idx}" ) logging.info(f"Completed episode {ep_idx}") return True except Exception as e: logging.error(f"Error processing episode {ep_idx}: {e}") logging.error(traceback.format_exc()) return False def main(): args = parse_args() setup_logging() # Setup paths dataset_path = args.dataset_path output_path = args.output_path # Find all task-level subdirectories (in case dataset_path is a parent directory) items = list_directory(dataset_path) task_subdirs = [ item for item in items if os.path.isdir(os.path.join(dataset_path, item)) ] for task_name in task_subdirs: dir_path = os.path.join(dataset_path, task_name) logging.info(f"Found task subdirectory: {dir_path}") data_base_path = resolve_path(dir_path, args.data_path) videos_base_path = resolve_path(dir_path, args.videos_path) meta_episodes_path = resolve_path(dir_path, args.meta_episodes_path) info_path = resolve_path(dir_path, args.info_path) # Setup tmp_dir tmp_dir = args.tmp_dir or tempfile.gettempdir() # Detect FPS if not specified fps = args.fps or detect_fps(info_path) # Auto-detect number of workers if not specified num_workers = args.num_workers or min(16, os.cpu_count() or 4) logging.info(f"Using {num_workers} worker threads") # Discover chunks and cameras data_chunks = discover_chunks(data_base_path) video_chunks = discover_chunks(videos_base_path) cameras = discover_cameras(video_chunks) # Detect image dimensions camera_name_to_img_shape = detect_image_dimensions(info_path, cameras) metadata = discover_metadata(meta_episodes_path) # Build episode and video lookups episode_lookup = build_episode_lookup(data_chunks) video_lookup = build_video_lookup(video_chunks, cameras) # depth_lookup = build_depth_lookup(depth_chunks, cameras) logging.info(f"Found {len(episode_lookup)} episodes") # Create output directory os.makedirs(output_path, exist_ok=True) # Configuration for episode processing config = { "lowdim_columns": args.lowdim_columns, "frame_index_col": args.frame_index_col, "episode_index_col": args.episode_index_col, "depth_columns": args.depth_columns, "tmp_dir": tmp_dir, } # Prepare episode processing tasks episode_jobs = [ ( ep_idx, ep_path, task_name, video_lookup, cameras, output_path, fps, metadata, camera_name_to_img_shape, config, ) for ep_idx, ep_path in episode_lookup.items() ] # Debug - process sequentially if args.debug: logging.info("Debug mode: processing episodes sequentially.") results = [] for job in episode_jobs: result = process_episode(job) results.append(result) else: # Process episodes in parallel using ThreadPoolExecutor with ThreadPoolExecutor(max_workers=num_workers) as executor: results = list(executor.map(process_episode, episode_jobs)) success_count = sum(1 for r in results if r) logging.info( f"Processed {success_count} episodes successfully out of {len(episode_jobs)}" ) if __name__ == "__main__": main()