"""LeRobot v3.0 dataset export engine. Reads episodes + frames from the database, writes the standard LeRobot directory layout: exports// ├── meta/ │ ├── info.json # schema, fps, features, path templates │ ├── stats.json # global normalization statistics │ ├── tasks.jsonl # task descriptions │ └── episodes/ # chunked Parquet episode index ├── data/ # chunked Parquet frames └── videos/ # per-camera MP4 shards """ import json import os from pathlib import Path from typing import Callable, Sequence from uuid import UUID import pandas as pd import pyarrow as pa import pyarrow.parquet as pq import numpy as np from app.config import settings from app.database import async_session from app.models import Episode, Frame, AnnotationSegment from sqlalchemy import select from sqlalchemy.orm import selectinload class DatasetExporter: """Export collected episodes to LeRobot v3.0 format.""" def __init__( self, job_id: UUID, dataset_name: str, episode_ids: Sequence[UUID], split_ratio: float = 0.9, ): self.job_id = job_id self.dataset_name = dataset_name self.episode_ids = episode_ids self.split_ratio = split_ratio self.output_dir = str(settings.storage_root / "exports" / dataset_name) async def run(self, progress_callback: Callable[[float], None] | None = None): """Execute the full export pipeline.""" os.makedirs(self.output_dir, exist_ok=True) # 1. Load all frames for selected episodes frames_df, episode_meta = await self._load_data() if frames_df.empty: raise ValueError("No frames found for selected episodes") total_eps = len(episode_meta) self._report(progress_callback, 5) # 2. Assign tasks task_map = await self._build_task_map(episode_meta) self._report(progress_callback, 10) # 3. Train/val split at episode level frames_df, episode_meta = self._assign_split(frames_df, episode_meta) self._report(progress_callback, 15) # 4. Write meta/ camera_names = self._detect_cameras(frames_df) features = self._build_features(frames_df, camera_names) fps = episode_meta[0].get("fps", 30) if episode_meta else 30 self._write_info_json(features, fps, total_eps, len(frames_df)) self._write_tasks_jsonl(task_map, episode_meta) self._report(progress_callback, 25) # 5. Write data/ Parquet shards await self._write_data_parquet(frames_df, progress_callback) self._report(progress_callback, 70) # 6. Write videos/ MP4 shards await self._write_videos(episode_meta, frames_df, camera_names, fps, progress_callback) self._report(progress_callback, 85) # 7. meta/episodes/ index — 由 lerobot_export / hdf5_export 流程 # 统一写为每个 episode 一个 episode_XXXXXX.parquet 文件，这里不再写 # (旧引擎曾写 chunk-000.parquet 顶层文件,与标准命名冲突)。 self._report(progress_callback, 90) # 8. Compute and write stats.json self._write_stats_json(frames_df, features) from app.storage import sync_tree_to_remote_async await sync_tree_to_remote_async( Path(self.output_dir), Path("exports") / self.dataset_name, ) self._report(progress_callback, 100) # ── Data loading ────────────────────────────────── async def _load_data(self) -> tuple[pd.DataFrame, list[dict]]: """Load frames, episode metadata, and annotations. Frames are read from the database first; if empty, fall back to reading the original LeRobot parquet files from the session directory. """ async with async_session() as session: frames_list = [] episode_meta = [] for ep_id in self.episode_ids: ep = await session.get(Episode, ep_id) if not ep: continue # ── Load annotations for this episode ── anno_q = ( select(AnnotationSegment) .options(selectinload(AnnotationSegment.keyframes)) .where(AnnotationSegment.episode_id == ep_id) .order_by(AnnotationSegment.sort_order, AnnotationSegment.start_frame_index) ) anno_segs = (await session.execute(anno_q)).scalars().all() # Build frame → (label, index) map frame_anno: dict[int, tuple[str | None, int]] = {} anno_defs = [] for idx, seg in enumerate(anno_segs): for fi in range(seg.start_frame_index, seg.end_frame_index + 1): if fi not in frame_anno: frame_anno[fi] = (seg.label, idx) kf_frames = [] for kf in seg.keyframes: kf_frames.append({ "frame_index": kf.frame_index, "event": kf.event or "", }) anno_defs.append({ "idx": idx, "label": seg.label, "start_frame": seg.start_frame_index, "end_frame": seg.end_frame_index, "frames": seg.end_frame_index - seg.start_frame_index + 1, "color": seg.color, "keyframes": kf_frames, }) episode_meta.append({ "episode_id": str(ep.id), "task_description": ep.task_description or "", "fps": ep.fps, "frame_count": ep.frame_count, "camera_names": ep.camera_names or [], "annotations": anno_defs, }) # ── Load frames: DB first, then parquet fallback ── q = ( select(Frame) .where(Frame.episode_id == ep_id) .order_by(Frame.frame_index) ) db_frames = (await session.execute(q)).scalars().all() if db_frames: # Path 1: frames stored in database for f in db_frames: row = self._build_frame_row(f, frame_anno) frames_list.append(row) else: # Path 2: frames in session parquet files — read directly parquet_frames = await self._load_frames_from_parquet( ep, frame_anno ) frames_list.extend(parquet_frames) df = pd.DataFrame(frames_list) return df, episode_meta def _build_frame_row(self, f: Frame, frame_anno: dict) -> dict: """Build a single frame row from a DB Frame record with annotations.""" anno_label, anno_idx = frame_anno.get(f.frame_index, (None, -1)) row = { "episode_id": str(f.episode_id), "frame_index": f.frame_index, "timestamp": f.timestamp, "reward": f.reward, "is_terminal": f.is_terminal, "is_truncated": f.is_truncated, "annotation": anno_label, "annotation_index": anno_idx, "annotation.language_persistent": anno_label, # 避开官方保留列名(结构化消息行) } obs = f.observation or {} if "state" in obs: row["observation.state"] = obs["state"] for k, v in obs.items(): if k != "state": row[f"observation.{k}"] = v act = f.action or {} if act: if "joint_positions" in act: row["action"] = act["joint_positions"] elif isinstance(act, list): row["action"] = act else: row["action"] = act if f.image_paths: for cam, path in f.image_paths.items(): row[f"image.{cam}"] = path return row async def _load_frames_from_parquet( self, ep: Episode, frame_anno: dict ) -> list[dict]: """Read frames from the original LeRobot parquet files in the session directory.""" from app.storage import find_session_dir_async frames_list = [] fps = ep.fps or 30 ep_id_str = str(ep.id) # Find the correct session directory for this episode session_dir = None if ep.session_id: async with async_session() as s: session_dir = await find_session_dir_async(str(ep.session_id), s) if session_dir is None: return frames_list # Scan for parquet files data_dir = None for d in session_dir.rglob("data"): if d.is_dir(): data_dir = d break if data_dir is None: for d in session_dir.rglob("*.parquet"): data_dir = d.parent break if data_dir is None: return frames_list # Read and annotate parquet files for parq_file in sorted(data_dir.rglob("*.parquet")): try: df = pd.read_parquet(parq_file) except Exception: continue # Skip meta parquets if not any(c for c in df.columns if "frame_index" in str(c).lower()): continue for _, row in df.iterrows(): fi = int(row.get("frame_index", len(frames_list))) anno_label, anno_idx = frame_anno.get(fi, (None, -1)) frame_row = { "episode_id": ep_id_str, "frame_index": fi, "timestamp": float(row.get("timestamp", fi / fps)), "reward": float(row.get("reward", 0.0)) if "reward" in df.columns else None, "is_terminal": bool(row.get("is_terminal", False)) if "is_terminal" in df.columns else False, "is_truncated": bool(row.get("is_truncated", False)) if "is_truncated" in df.columns else False, "annotation": anno_label, "annotation_index": anno_idx, "annotation.language_persistent": anno_label, # 避开官方保留列名(结构化消息行) } # Copy observation/action columns for col in df.columns: if col.startswith("observation."): val = row[col] if isinstance(val, (list, pd.Series)): val = list(val) frame_row[col] = val elif col == "action" or col.startswith("action"): val = row[col] if isinstance(val, (list, pd.Series)): val = list(val) elif hasattr(val, 'item'): val = val.item() if hasattr(val, 'size') and val.size == 1 else list(val) frame_row["action"] = val elif col.startswith("episode_index"): pass # handled above elif col.startswith("task_index"): pass # handled by task assignment later # Copy image paths for col in df.columns: if col.startswith("image.") or col.startswith("video."): frame_row[col] = str(row[col]) if not pd.isna(row[col]) else None frames_list.append(frame_row) return frames_list # ── Tasks ───────────────────────────────────────── async def _build_task_map(self, episode_meta: list[dict]) -> dict[int, str]: """Build a map of task_id → description.""" tasks: dict[str, int] = {} task_map: dict[int, str] = {} idx = 0 for ep in episode_meta: desc = ep.get("task_description", "").strip() if desc and desc not in tasks: tasks[desc] = idx task_map[idx] = desc idx += 1 # Assign task_id to each episode for ep in episode_meta: desc = ep.get("task_description", "").strip() ep["task_id"] = tasks.get(desc, -1) if desc else -1 return task_map # ── Split ───────────────────────────────────────── def _assign_split(self, df: pd.DataFrame, episode_meta: list[dict]) -> tuple[pd.DataFrame, list[dict]]: """Assign train/val split at episode level.""" rng = np.random.default_rng(42) ep_list = list(set(df["episode_id"])) rng.shuffle(ep_list) n_train = max(1, int(len(ep_list) * self.split_ratio)) train_eps = set(ep_list[:n_train]) df["split"] = df["episode_id"].apply(lambda e: "train" if e in train_eps else "val") for ep in episode_meta: ep["split"] = "train" if ep["episode_id"] in train_eps else "val" return df, episode_meta # ── Features detection ──────────────────────────── def _detect_cameras(self, df: pd.DataFrame) -> list[str]: """Detect camera names from the dataframe columns.""" cameras = set() for col in df.columns: if col.startswith("image."): cameras.add(col.replace("image.", "")) return sorted(cameras) def _build_features(self, df: pd.DataFrame, camera_names: list[str]) -> dict: """Build the features dict for info.json.""" features = {} for col in df.columns: if col.startswith("observation.") and col != "observation.state": val = df[col].dropna().iloc[0] if not df[col].dropna().empty else None if val is not None: if isinstance(val, (list, np.ndarray)): features[col] = {"dtype": "float32", "shape": [len(val)]} else: features[col] = {"dtype": "float32", "shape": [1]} elif col == "observation.state": val = df[col].dropna().iloc[0] if not df[col].dropna().empty else None if val is not None and isinstance(val, (list, np.ndarray)): features[col] = {"dtype": "float32", "shape": [len(val)]} elif col == "action": val = df[col].dropna().iloc[0] if not df[col].dropna().empty else None if val is not None and isinstance(val, (list, np.ndarray)): features[col] = {"dtype": "float32", "shape": [len(val)]} for cam in camera_names: features[f"observation.images.{cam}"] = { "dtype": "video", "shape": [480, 640, 3], "info": {"camera": cam}, } # Annotation features if "annotation" in df.columns: features["annotation"] = {"dtype": "string", "shape": [1]} if "annotation_index" in df.columns: features["annotation_index"] = {"dtype": "int64", "shape": [1]} if "annotation.language_persistent" in df.columns: features["annotation.language_persistent"] = {"dtype": "string", "shape": [1]} return features # ── meta/ writers ───────────────────────────────── def _write_info_json(self, features: dict, fps: int, total_episodes: int, total_frames: int): """Write meta/info.json.""" info = { "codebase_version": "v3.0", "robot_type": "unknown", "fps": fps, "features": features, "total_episodes": total_episodes, "total_frames": total_frames, "splits": { "train": "data/train-*.parquet", "val": "data/val-*.parquet", }, "data_path": "data/{split}-{shard_idx:03d}.parquet", "video_path": "videos/chunk-{chunk_idx:03d}/{video_key}/file-{file_idx:03d}.mp4", "video_keys": [ f"observation.images.{cam}" for cam in self._detect_cameras(pd.DataFrame()) ] or [], } # Add detected camera keys if features: cam_keys = [k for k in features if k.startswith("observation.images.")] if cam_keys: info["video_keys"] = cam_keys path = Path(self.output_dir) / "meta" path.mkdir(parents=True, exist_ok=True) with open(path / "info.json", "w") as f: json.dump(info, f, indent=2) def _write_tasks_jsonl(self, task_map: dict[int, str], episode_meta: list[dict] | None = None): """Write meta/tasks.jsonl with optional annotation definitions.""" path = Path(self.output_dir) / "meta" path.mkdir(parents=True, exist_ok=True) # Collect annotation definitions per task task_annotations: dict[int, list[dict]] = {} if episode_meta: for ep in episode_meta: tid = ep.get("task_id", -1) annos = ep.get("annotations", []) if annos and tid >= 0 and tid not in task_annotations: task_annotations[tid] = annos with open(path / "tasks.jsonl", "w") as f: for tid, desc in sorted(task_map.items()): entry: dict = {"task_id": tid, "description": desc} if tid in task_annotations: entry["annotations"] = task_annotations[tid] entry["annotation_schema"] = { "column": "annotation", "dtype": "string", "description": "Per-frame action label", } f.write(json.dumps(entry, ensure_ascii=False) + "\n") # ── data/ Parquet ───────────────────────────────── async def _write_data_parquet(self, df: pd.DataFrame, progress_callback=None): """Write data/ as chunked Parquet files, split by train/val.""" data_path = Path(self.output_dir) / "data" data_path.mkdir(parents=True, exist_ok=True) for split_name in ["train", "val"]: split_df = df[df["split"] == split_name] if split_df.empty: continue shard_idx = 0 for i in range(0, len(split_df), settings.EXPORT_SHARD_SIZE): chunk = split_df.iloc[i : i + settings.EXPORT_SHARD_SIZE] t = pa.Table.from_pandas(chunk) pq.write_table(t, str(data_path / f"{split_name}-{shard_idx:03d}.parquet")) shard_idx += 1 if progress_callback: progress_callback(25 + int(45 * (i / len(split_df)) * 0.5)) # ── videos/ MP4 ─────────────────────────────────── async def _write_videos( self, episode_meta: list[dict], df: pd.DataFrame, camera_names: list[str], fps: int, progress_callback=None, ): """Write videos/ as per-camera MP4 shards from saved frame images. Uses OpenCV to encode JPEG frames into MP4. """ if not camera_names: return try: import cv2 except ImportError: print("[Export] opencv not available, skipping video encoding") return video_path = Path(self.output_dir) / "videos" video_path.mkdir(parents=True, exist_ok=True) for cam in camera_names: cam_dir = video_path / "chunk-000" / f"observation.images.{cam}" cam_dir.mkdir(parents=True, exist_ok=True) image_col = f"image.{cam}" if image_col not in df.columns: continue # Group frames by episode and encode for ep in episode_meta: ep_df = df[df["episode_id"] == ep["episode_id"]].sort_values("frame_index") images = ep_df[image_col].dropna() if images.empty: continue # Determine frame size from first image first_img_path = Path(settings.STORAGE_DIR) / images.iloc[0] if not first_img_path.exists(): continue sample = cv2.imread(str(first_img_path)) if sample is None: continue h, w = sample.shape[:2] fourcc = cv2.VideoWriter_fourcc(*"mp4v") out_file = str(cam_dir / f"episode_{ep['episode_id'][:8]}.mp4") writer = cv2.VideoWriter(out_file, fourcc, fps, (w, h)) for _, row in ep_df.iterrows(): img_rel = row.get(image_col) if not img_rel or pd.isna(img_rel): # Write duplicate frame to maintain timing if writer.isOpened(): writer.write(sample) continue img_abs = Path(settings.STORAGE_DIR) / img_rel if img_abs.exists(): frame = cv2.imread(str(img_abs)) if frame is not None: # Resize if needed if frame.shape[:2] != (h, w): frame = cv2.resize(frame, (w, h)) writer.write(frame) sample = frame # fallback for missing frames writer.release() # ── stats.json ──────────────────────────────────── def _write_stats_json(self, df: pd.DataFrame, features: dict): """Compute and write meta/stats.json with global normalization statistics.""" stats = {} for feat_name, feat_info in features.items(): if feat_name.startswith("observation.images."): continue # skip video features if feat_name not in df.columns: continue series = df[feat_name].dropna() if series.empty: continue # Flatten array values values = [] for v in series: if isinstance(v, (list, np.ndarray)): values.append(np.array(v, dtype=np.float32)) elif isinstance(v, (int, float)): values.append(np.array([v], dtype=np.float32)) if not values: continue arr = np.stack(values) stats[feat_name] = { "mean": arr.mean(axis=0).tolist(), "std": arr.std(axis=0).tolist(), "min": arr.min(axis=0).tolist(), "max": arr.max(axis=0).tolist(), "count": len(values), } path = Path(self.output_dir) / "meta" with open(path / "stats.json", "w") as f: json.dump(stats, f, indent=2) def _report(self, callback, pct): if callback: callback(pct)