validate_converted.py

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