""" Convert the T-Rex LeRobot v3.0 dataset to LeRobot v2.1 layout for DreamZero. LeRobot v3 packs many episodes into shared parquet/video files: data/chunk-XXX/file-XXX.parquet (rows of many episodes) videos/{video_key}/chunk-XXX/file-XXX.mp4 (concatenated episodes) meta/episodes/chunk-XXX/file-XXX.parquet (episode metadata) meta/tasks.parquet DreamZero's loader (groot/vla/data/dataset/lerobot.py) expects v2 layout: data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4 meta/episodes.jsonl, meta/tasks.jsonl, meta/info.json RGB cameras are re-encoded to 320x180. Tactile videos keep native resolution (raw 320x240, deform 240x240) and are re-encoded with libx264 for smaller size. Usage: python scripts/data/convert_trex_v3_to_v2.py --phase data python scripts/data/convert_trex_v3_to_v2.py --phase videos python scripts/data/convert_trex_v3_to_v2.py --phase videos --include-tactile python scripts/data/convert_trex_v3_to_v2.py --phase meta --include-tactile python scripts/data/convert_trex_v3_to_v2.py --phase verify --include-tactile All phases are resumable: existing valid outputs are skipped. """ from __future__ import annotations import argparse import json import logging import subprocess from concurrent.futures import ProcessPoolExecutor, as_completed from pathlib import Path import numpy as np import pandas as pd from tqdm import tqdm logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") log = logging.getLogger(__name__) DEFAULT_SRC = Path("/scratch1/home/zhicao/dreamzero/data/trex_dataset") DEFAULT_DST = Path("/scratch1/home/zhicao/dreamzero/data/trex_datasetv2") SRC = DEFAULT_SRC DST = DEFAULT_DST RGB_VIDEO_KEYS = [ "observation.images.head_left", "observation.images.left_wrist", "observation.images.right_wrist", ] RGB_OUT_W, RGB_OUT_H = 320, 180 FPS = 30 CHUNKS_SIZE = 1000 RGB_CRF = 23 TACTILE_CRF = 28 def load_src_info() -> dict: return json.loads((SRC / "meta" / "info.json").read_text()) def get_tactile_video_keys(src_info: dict | None = None) -> list[str]: src_info = src_info or load_src_info() return sorted( k for k, v in src_info["features"].items() if v.get("dtype") == "video" and "tactile" in k ) def get_output_size(video_key: str, feature: dict) -> tuple[int, int]: """Return (width, height) for ffmpeg scale filter.""" if video_key in RGB_VIDEO_KEYS: return RGB_OUT_W, RGB_OUT_H shape = feature.get("shape", []) if len(shape) >= 2: height, width = int(shape[0]), int(shape[1]) return width, height info = feature.get("info", {}) return int(info["video.width"]), int(info["video.height"]) def get_crf(video_key: str) -> int: return RGB_CRF if video_key in RGB_VIDEO_KEYS else TACTILE_CRF def resolve_video_keys( include_tactile: bool, tactile_only: bool, explicit_keys: list[str] | None, ) -> list[str]: if explicit_keys: return explicit_keys if tactile_only: return get_tactile_video_keys() keys = list(RGB_VIDEO_KEYS) if include_tactile: keys.extend(get_tactile_video_keys()) return keys def load_episode_meta() -> pd.DataFrame: files = sorted(SRC.glob("meta/episodes/chunk-*/file-*.parquet")) if not files: raise FileNotFoundError(f"No episode metadata under {SRC / 'meta/episodes'}") df = pd.concat([pd.read_parquet(f) for f in files], ignore_index=True) return df.sort_values("episode_index").reset_index(drop=True) def load_task_map() -> dict[int, str]: t = pd.read_parquet(SRC / "meta" / "tasks.parquet") return {int(row.task_index): str(idx) for idx, row in t.iterrows()} def ep_parquet_path(ep_idx: int) -> Path: return DST / f"data/chunk-{ep_idx // CHUNKS_SIZE:03d}/episode_{ep_idx:06d}.parquet" def ep_video_path(ep_idx: int, video_key: str) -> Path: return DST / f"videos/chunk-{ep_idx // CHUNKS_SIZE:03d}/{video_key}/episode_{ep_idx:06d}.mp4" def build_video_features(src_info: dict, video_keys: list[str]) -> dict: features: dict = {} for k, v in src_info["features"].items(): if v.get("dtype") != "video" or k not in video_keys: continue v = dict(v) out_w, out_h = get_output_size(k, v) v["shape"] = [out_h, out_w, 3] info_blk = dict(v.get("info", {})) info_blk.update( { "video.height": out_h, "video.width": out_w, "video.codec": "h264", "video.pix_fmt": "yuv420p", "video.fps": FPS, "video.channels": 3, "has_audio": False, } ) v["info"] = info_blk features[k] = v return features def write_info_json( src_info: dict, video_keys: list[str], *, preserve_existing_features: bool = False, ) -> None: meta_dir = DST / "meta" meta_dir.mkdir(parents=True, exist_ok=True) existing = {} info_path = meta_dir / "info.json" if preserve_existing_features and info_path.exists(): existing = json.loads(info_path.read_text()) features = dict(existing.get("features", {})) for k, v in src_info["features"].items(): if v.get("dtype") != "video": features[k] = v features.update(build_video_features(src_info, video_keys)) features["annotation.task"] = {"dtype": "string", "shape": [1], "names": None} total_episodes = int(src_info["total_episodes"]) all_video_keys = [k for k, v in features.items() if v.get("dtype") == "video"] info = { "codebase_version": "v2.1", "robot_type": src_info.get("robot_type", "dexmate_vega1_and_sharpa_wave"), "total_episodes": total_episodes, "total_frames": int(src_info["total_frames"]), "total_tasks": int(src_info["total_tasks"]), "total_videos": total_episodes * len(all_video_keys), "total_chunks": (total_episodes + CHUNKS_SIZE - 1) // CHUNKS_SIZE, "chunks_size": CHUNKS_SIZE, "fps": FPS, "splits": {"train": f"0:{total_episodes}"}, "data_path": "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet", "video_path": "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4", "features": features, } with open(info_path, "w") as f: json.dump(info, f, indent=4) log.info( "Wrote meta/info.json with %d video keys (%d total videos)", len(all_video_keys), info["total_videos"], ) # --------------------------------------------------------------------------- # Phase: data (parquets + meta) # --------------------------------------------------------------------------- def convert_data( ep_meta: pd.DataFrame, task_map: dict[int, str], video_keys: list[str], ) -> None: src_info = load_src_info() n_done = 0 groups = ep_meta.groupby(["data/chunk_index", "data/file_index"]) for (chunk_idx, file_idx), eps in tqdm(groups, desc="Converting data files"): src_pq = SRC / f"data/chunk-{int(chunk_idx):03d}/file-{int(file_idx):03d}.parquet" if not src_pq.exists(): log.warning("Missing source parquet: %s", src_pq) continue if all(ep_parquet_path(int(r.episode_index)).exists() for r in eps.itertuples()): n_done += len(eps) continue df = pd.read_parquet(src_pq) for r in eps.itertuples(): ep_idx = int(r.episode_index) out = ep_parquet_path(ep_idx) if out.exists(): n_done += 1 continue ep_df = df[df["episode_index"] == ep_idx].copy() assert len(ep_df) == int(r.length), ( f"episode {ep_idx}: rows {len(ep_df)} != meta length {r.length}" ) task_texts = [str(t) for t in r.tasks] ep_df["annotation.task"] = task_texts[0] if task_texts else "" out.parent.mkdir(parents=True, exist_ok=True) ep_df.to_parquet(out, index=False) n_done += 1 log.info("Data phase done: %d episode parquets", n_done) meta_dir = DST / "meta" meta_dir.mkdir(parents=True, exist_ok=True) with open(meta_dir / "tasks.jsonl", "w") as f: for idx in sorted(task_map): f.write(json.dumps({"task_index": idx, "task": task_map[idx]}) + "\n") with open(meta_dir / "episodes.jsonl", "w") as f: for r in ep_meta.itertuples(): f.write( json.dumps( { "episode_index": int(r.episode_index), "tasks": [str(t) for t in r.tasks], "length": int(r.length), } ) + "\n" ) write_info_json(src_info, video_keys, preserve_existing_features=False) log.info("Wrote meta/episodes.jsonl, meta/tasks.jsonl") def update_meta(video_keys: list[str]) -> None: src_info = load_src_info() write_info_json(src_info, video_keys, preserve_existing_features=True) # --------------------------------------------------------------------------- # Phase: videos # --------------------------------------------------------------------------- def _cut_one(job: tuple) -> tuple[int, str, bool, str]: ep_idx, video_key, src_mp4, from_ts, n_frames, out_path, out_w, out_h, crf = job out = Path(out_path) out.parent.mkdir(parents=True, exist_ok=True) tmp = out.with_suffix(".tmp.mp4") ss = max(0.0, from_ts - 0.5 / FPS) cmd = [ "ffmpeg", "-y", "-loglevel", "error", "-ss", f"{ss:.6f}", "-i", src_mp4, "-frames:v", str(n_frames), "-vf", f"scale={out_w}:{out_h}", "-c:v", "libx264", "-preset", "veryfast", "-crf", str(crf), "-pix_fmt", "yuv420p", "-movflags", "+faststart", "-an", "-threads", "2", str(tmp), ] try: res = subprocess.run(cmd, capture_output=True, text=True, timeout=600) if res.returncode != 0: tmp.unlink(missing_ok=True) return ep_idx, video_key, False, res.stderr[-500:] tmp.rename(out) return ep_idx, video_key, True, "" except Exception as e: # noqa: BLE001 tmp.unlink(missing_ok=True) return ep_idx, video_key, False, str(e) def convert_videos(ep_meta: pd.DataFrame, video_keys: list[str], workers: int) -> None: src_info = load_src_info() jobs = [] missing_src = set() for vk in video_keys: feature = src_info["features"][vk] out_w, out_h = get_output_size(vk, feature) crf = get_crf(vk) for r in ep_meta.itertuples(): ep_idx = int(r.episode_index) out = ep_video_path(ep_idx, vk) if out.exists(): continue chunk_i = int(ep_meta.loc[r.Index, f"videos/{vk}/chunk_index"]) file_i = int(ep_meta.loc[r.Index, f"videos/{vk}/file_index"]) from_ts = float(ep_meta.loc[r.Index, f"videos/{vk}/from_timestamp"]) src_mp4 = SRC / f"videos/{vk}/chunk-{chunk_i:03d}/file-{file_i:03d}.mp4" if not src_mp4.exists(): missing_src.add(str(src_mp4)) continue jobs.append( ( ep_idx, vk, str(src_mp4), from_ts, int(r.length), str(out), out_w, out_h, crf, ) ) if missing_src: log.warning( "%d source videos missing (not yet downloaded?), e.g. %s", len(missing_src), sorted(missing_src)[0], ) log.info("Cutting %d episode videos with %d workers", len(jobs), workers) failures = [] with ProcessPoolExecutor(max_workers=workers) as pool: futs = [pool.submit(_cut_one, j) for j in jobs] for fut in tqdm(as_completed(futs), total=len(futs), desc="Cutting videos"): ep_idx, vk, ok, err = fut.result() if not ok: failures.append((ep_idx, vk, err)) if failures: log.error("%d failures, first: %s", len(failures), failures[0]) else: log.info("Video phase done, no failures") # --------------------------------------------------------------------------- # Phase: verify # --------------------------------------------------------------------------- def _probe_frames(path: Path) -> int: res = subprocess.run( [ "ffprobe", "-v", "error", "-count_frames", "-select_streams", "v:0", "-show_entries", "stream=nb_read_frames", "-of", "csv=p=0", str(path), ], capture_output=True, text=True, timeout=120, ) return int(res.stdout.strip()) def verify(ep_meta: pd.DataFrame, video_keys: list[str], n_samples: int) -> None: rng = np.random.default_rng(0) total = len(ep_meta) missing_pq = [ int(r.episode_index) for r in ep_meta.itertuples() if not ep_parquet_path(int(r.episode_index)).exists() ] log.info("Parquets: %d/%d present", total - len(missing_pq), total) for vk in video_keys: missing = [ int(r.episode_index) for r in ep_meta.itertuples() if not ep_video_path(int(r.episode_index), vk).exists() ] log.info("Videos [%s]: %d/%d present", vk, total - len(missing), total) sample = rng.choice(total, size=min(n_samples, total), replace=False) for ep_idx in sample: ep_idx = int(ep_idx) row = ep_meta[ep_meta["episode_index"] == ep_idx].iloc[0] length = int(row["length"]) pq = ep_parquet_path(ep_idx) if pq.exists(): df = pd.read_parquet(pq) assert len(df) == length, f"ep {ep_idx}: parquet {len(df)} != {length}" assert np.asarray(df["action"].iloc[0]).shape == (58,) assert df["annotation.task"].iloc[0] == str(row["tasks"][0]) for vk in video_keys: vp = ep_video_path(ep_idx, vk) if vp.exists(): n = _probe_frames(vp) assert n == length, f"ep {ep_idx} {vk}: video {n} frames != {length}" log.info("ep %06d OK (length=%d)", ep_idx, length) log.info("Verification passed on %d sampled episodes", len(sample)) def main() -> None: global SRC, DST parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--phase", choices=["data", "videos", "meta", "verify"], required=True, ) parser.add_argument("--src", type=Path, default=DEFAULT_SRC, help="v3 dataset root") parser.add_argument("--dst", type=Path, default=DEFAULT_DST, help="v2 output root") parser.add_argument( "--include-tactile", action="store_true", help="Include all 20 tactile video streams", ) parser.add_argument( "--tactile-only", action="store_true", help="Convert/update only tactile video streams (skip RGB)", ) parser.add_argument( "--video-keys", nargs="+", default=None, help="Explicit video keys to convert (overrides --include-tactile default set)", ) parser.add_argument("--workers", type=int, default=16) parser.add_argument("--verify-samples", type=int, default=20) args = parser.parse_args() SRC = args.src DST = args.dst ep_meta = load_episode_meta() log.info("Loaded %d episodes from v3 metadata", len(ep_meta)) video_keys = resolve_video_keys(args.include_tactile, args.tactile_only, args.video_keys) log.info("Video keys (%d): %s", len(video_keys), ", ".join(video_keys)) if args.phase == "data": convert_data(ep_meta, load_task_map(), video_keys) elif args.phase == "meta": update_meta(video_keys) elif args.phase == "videos": convert_videos(ep_meta, video_keys, args.workers) elif args.phase == "verify": verify(ep_meta, video_keys, args.verify_samples) if __name__ == "__main__": main()