#!/usr/bin/env python3 """Exhaustively validate converted LeRobot v3.0 tabular/video metadata.""" from __future__ import annotations import json import shutil import subprocess import sys from dataclasses import dataclass from pathlib import Path from typing import Any import numpy as np import pyarrow as pa import pyarrow.parquet as pq FPS = 30 VECTOR_COLUMNS = ( "observation.state", "observation.velocity", "observation.effort", "action", ) CAMERAS = ( "observation.images.cam_high", "observation.images.cam_left_wrist", "observation.images.cam_right_wrist", ) @dataclass(frozen=True) class DatasetSpec: name: str episodes: int frames: int task: str SPECS = ( DatasetSpec( name="table_clean", episodes=100, frames=89_469, task=( "pick up the crumpled paper and small blocks from the tabletop, place them " "into the tray, then use the cloth to wipe the brown stain on the table." ), ), DatasetSpec( name="put_mango", episodes=100, frames=31_000, task="put the mango on the plate", ), ) def fail(name: str, message: str) -> None: raise ValueError(f"{name}: {message}") def read_parquet_tree(root: Path) -> pa.Table: files = sorted(root.rglob("*.parquet")) if not files: raise FileNotFoundError(f"No parquet files under {root}") return pa.concat_tables([pq.read_table(path) for path in files]) def scalar_array(table: pa.Table, column: str) -> np.ndarray: return table[column].combine_chunks().to_numpy(zero_copy_only=False) def list_values(table: pa.Table, column: str, expected_width: int) -> np.ndarray: values = table[column].combine_chunks() lengths = np.diff(values.offsets.to_numpy()) if not np.all(lengths == expected_width): raise ValueError( f"{column}: expected every row to have width {expected_width}, " f"got widths {np.unique(lengths).tolist()}" ) return values.values.to_numpy(zero_copy_only=False).reshape(-1, expected_width) def find_ffprobe() -> Path: candidate = shutil.which("ffprobe") if candidate: return Path(candidate) bundled = Path(sys.executable).resolve().parent / "Library" / "bin" / "ffprobe.exe" if bundled.is_file(): return bundled raise FileNotFoundError("ffprobe was not found on PATH or in the active environment") def video_frame_count(ffprobe: Path, path: Path) -> int: result = subprocess.run( [ str(ffprobe), "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=nb_frames", "-of", "default=nokey=1:noprint_wrappers=1", str(path), ], check=True, capture_output=True, text=True, encoding="utf-8", ) return int(result.stdout.strip()) def validate_dataset(root: Path, spec: DatasetSpec, ffprobe: Path) -> dict[str, Any]: info = json.loads((root / "meta" / "info.json").read_text(encoding="utf-8")) if info["codebase_version"] != "v3.0": fail(spec.name, f"expected codebase_version v3.0, got {info['codebase_version']}") if info["total_episodes"] != spec.episodes or info["total_frames"] != spec.frames: fail(spec.name, "info.json episode/frame totals do not match the expected values") if info["fps"] != FPS: fail(spec.name, f"expected {FPS} FPS, got {info['fps']}") data = read_parquet_tree(root / "data") episodes_meta = read_parquet_tree(root / "meta" / "episodes") if data.num_rows != spec.frames or episodes_meta.num_rows != spec.episodes: fail(spec.name, "Parquet row totals do not match info.json") indices = scalar_array(data, "index") episode_indices = scalar_array(data, "episode_index") frame_indices = scalar_array(data, "frame_index") timestamps = scalar_array(data, "timestamp") task_indices = scalar_array(data, "task_index") if not np.array_equal(indices, np.arange(spec.frames)): fail(spec.name, "global index is not contiguous from zero") if not np.all(task_indices == 0): fail(spec.name, "unexpected task_index value") if not np.isfinite(timestamps).all(): fail(spec.name, "timestamp contains NaN or Inf") numeric_summaries: dict[str, Any] = {} for column in VECTOR_COLUMNS: values = list_values(data, column, expected_width=14) if not np.isfinite(values).all(): fail(spec.name, f"{column} contains NaN or Inf") numeric_summaries[column] = { "shape": list(values.shape), "min": float(values.min()), "max": float(values.max()), } episode_ids = scalar_array(episodes_meta, "episode_index") starts = scalar_array(episodes_meta, "dataset_from_index") stops = scalar_array(episodes_meta, "dataset_to_index") lengths = scalar_array(episodes_meta, "length") if not np.array_equal(episode_ids, np.arange(spec.episodes)): fail(spec.name, "episode metadata index is not contiguous from zero") if starts[0] != 0 or stops[-1] != spec.frames: fail(spec.name, "episode metadata does not cover the full dataset") if not np.array_equal(starts[1:], stops[:-1]) or not np.array_equal(stops - starts, lengths): fail(spec.name, "episode metadata has gaps, overlaps, or invalid lengths") max_timestamp_error_s = 0.0 for episode_id, start, stop, length in zip(episode_ids, starts, stops, lengths, strict=True): selection = slice(int(start), int(stop)) if not np.all(episode_indices[selection] == episode_id): fail(spec.name, f"episode_index mismatch in episode {episode_id}") if not np.array_equal(frame_indices[selection], np.arange(length)): fail(spec.name, f"frame_index mismatch in episode {episode_id}") expected_timestamps = np.arange(length, dtype=np.float64) / FPS error = float(np.max(np.abs(timestamps[selection] - expected_timestamps))) max_timestamp_error_s = max(max_timestamp_error_s, error) if error > 2e-6: fail(spec.name, f"timestamp cadence mismatch in episode {episode_id}: {error}s") tasks = pq.read_table(root / "meta" / "tasks.parquet").to_pydict() if tasks != {"task_index": [0], "task": [spec.task]}: fail(spec.name, f"unexpected task metadata: {tasks}") episode_tasks = episodes_meta["tasks"].to_pylist() if any(value != [spec.task] for value in episode_tasks): fail(spec.name, "episode task labels are inconsistent") video_summary: dict[str, Any] = {} for camera in CAMERAS: files = sorted((root / "videos" / camera).rglob("*.mp4")) counts = [video_frame_count(ffprobe, path) for path in files] if sum(counts) != spec.frames: fail(spec.name, f"{camera} has {sum(counts)} video frames, expected {spec.frames}") from_col = scalar_array(episodes_meta, f"videos/{camera}/from_timestamp") to_col = scalar_array(episodes_meta, f"videos/{camera}/to_timestamp") duration_error = float(np.max(np.abs((to_col - from_col) - (lengths / FPS)))) if duration_error > 2e-9: fail(spec.name, f"{camera} episode duration mismatch: {duration_error}s") video_summary[camera] = { "shards": len(files), "frames_per_shard": counts, "total_frames": sum(counts), "max_episode_duration_error_s": duration_error, } return { "name": spec.name, "root": str(root), "codebase_version": info["codebase_version"], "episodes": spec.episodes, "frames": spec.frames, "fps": FPS, "task": spec.task, "max_timestamp_error_s": max_timestamp_error_s, "numeric": numeric_summaries, "videos": video_summary, } def main() -> None: output_root = Path("data/converted").resolve() ffprobe = find_ffprobe() results = [] for spec in SPECS: result = validate_dataset(output_root / spec.name, spec, ffprobe) results.append(result) print( f"[{spec.name}] OK: {spec.episodes} episodes, {spec.frames} rows, " f"3 cameras, all numeric values finite", flush=True, ) report = { "status": "passed", "ffprobe": str(ffprobe), "totals": { "datasets": len(results), "episodes": sum(item["episodes"] for item in results), "frames": sum(item["frames"] for item in results), }, "datasets": results, } report_path = output_root / "full_validation.json" report_path.write_text( json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) print(f"Validation report: {report_path}") if __name__ == "__main__": main()