-
Edit Dataset
-
- This page is under construction.
-
+
+
+
+
navigate("/")}
+ className="h-8 gap-1.5 px-2 text-gray-400 hover:bg-gray-900 hover:text-white"
+ >
+
+
+
+
+
+
+
+
+
+
+ {localDatasets.map((d) => (
+
+ {d.repo_id}
+
+ ))}
+
+
+
+ {/* Picking a dataset now lands here rather than on /upload, so the
+ upload + delete flow hangs off the page you browse from. */}
+ {repoId && (
+
+ navigate("/upload", {
+ state: {
+ datasetInfo: {
+ dataset_repo_id: repoId,
+ source: selectedSource ?? "local",
+ },
+ },
+ })
+ }
+ className="h-9 gap-1.5 border-gray-800 bg-gray-950 text-xs text-gray-300 hover:bg-gray-900"
+ >
+
+ )}
+
+
+
+ {!repoId && (
+
+
+
Select a dataset to browse its episodes.
+ {!datasetsLoading && localDatasets.length === 0 && (
+
+ No datasets in the local LeRobot cache. Record one, or download it from the
+ Hub first.
+
+ )}
+
+ )}
+
+ {repoId && loading && (
+
+
+ )}
+
+ {repoId && error && (
+
+
+
+
Could not read this dataset
+
{error}
+
+
+ )}
+
+ {repoId && data?.success && (
+ <>
+
+
+
+
+
+
+
+ {detailLoading && !detail && (
+
+
+ )}
+ {detailError && (
+
+ {detailError}
+
+ )}
+ {detail && (
+
+ )}
+
+
+
+ )}
+
);
};
diff --git a/frontend/src/pages/Landing.tsx b/frontend/src/pages/Landing.tsx
index 33832569..b80b8470 100644
--- a/frontend/src/pages/Landing.tsx
+++ b/frontend/src/pages/Landing.tsx
@@ -97,15 +97,11 @@ const Landing = () => {
};
const handlePickExisting = (item: DatasetItem) => {
+ // Anything with files on disk opens in the local episode browser; upload and
+ // delete are reachable from there. A Hub-only dataset has no videos here to
+ // decode, so it still goes out to the visualize_dataset Space.
if (item.source === "local" || item.source === "both") {
- navigate("/upload", {
- state: {
- datasetInfo: {
- dataset_repo_id: item.repo_id,
- source: item.source,
- },
- },
- });
+ navigate(`/edit-dataset?dataset=${encodeURIComponent(item.repo_id)}`);
return;
}
openHubViewer(item.repo_id, item.private);
diff --git a/lelab/datasets.py b/lelab/datasets.py
index b34c85c3..715ce8be 100644
--- a/lelab/datasets.py
+++ b/lelab/datasets.py
@@ -12,21 +12,22 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import base64
import logging
-import os
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from huggingface_hub.errors import HfHubHTTPError
+from . import episode_media
from .utils.hf_auth import cached_whoami, shared_hf_api
logger = logging.getLogger(__name__)
def _lerobot_cache_root() -> Path:
- return Path(os.environ.get("HF_LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser()
+ return episode_media.lerobot_cache_root()
def _is_dataset_dir(path: Path) -> bool:
@@ -160,3 +161,184 @@ def list_all_datasets() -> list[dict[str, Any]]:
out = list(merged.values())
out.sort(key=lambda d: d["last_modified"] or "", reverse=True)
return out
+
+
+# ── episode browsing ────────────────────────────────────────────────────────
+#
+# Everything below reads a *local* dataset only. Browsing a Hub-only dataset
+# would mean pulling its videos down first, so the frontend offers the episode
+# browser for entries whose `source` is "local" or "both" and keeps sending
+# Hub-only ones to the visualize_dataset Space.
+
+
+def handle_list_episodes(repo_id: str) -> dict[str, Any]:
+ """Episode listing for one local dataset: per-episode length, duration, tasks."""
+ try:
+ dataset_dir = episode_media.resolve_dataset_dir(repo_id)
+ info = episode_media.read_info(dataset_dir)
+ cameras = episode_media.list_cameras(dataset_dir)
+ rows = episode_media.read_episode_index(dataset_dir)
+ except episode_media.DatasetNotFoundError as e:
+ return {"success": False, "message": str(e)}
+ except Exception as e:
+ logger.warning(f"Could not list episodes for {repo_id}: {e}")
+ return {"success": False, "message": f"Could not read dataset: {e}"}
+
+ return {
+ "success": True,
+ "repo_id": repo_id,
+ "fps": info.get("fps"),
+ "robot_type": info.get("robot_type"),
+ "total_episodes": info.get("total_episodes", len(rows)),
+ "total_frames": info.get("total_frames"),
+ "cameras": cameras,
+ "episodes": [
+ {
+ "episode_index": r.episode_idx,
+ "length": r.length,
+ "duration_s": r.duration_s,
+ "tasks": list(r.tasks),
+ }
+ for r in rows
+ ],
+ }
+
+
+def handle_episode_detail(repo_id: str, episode_index: int) -> dict[str, Any]:
+ """One episode: its per-camera video windows, plus neighbours for paging.
+
+ Each camera reports its own ``from_timestamp``. They agree in practice, but
+ the player must not assume that — a secondary camera is synced by mapping
+ through *its own* offset, not by copying the driver's ``currentTime``.
+ """
+ try:
+ dataset_dir = episode_media.resolve_dataset_dir(repo_id)
+ info = episode_media.read_info(dataset_dir)
+ rows = episode_media.read_episode_index(dataset_dir)
+ except episode_media.DatasetNotFoundError as e:
+ return {"success": False, "message": str(e)}
+ except Exception as e:
+ logger.warning(f"Could not read episode {episode_index} of {repo_id}: {e}")
+ return {"success": False, "message": f"Could not read dataset: {e}"}
+
+ row = next((r for r in rows if r.episode_idx == episode_index), None)
+ if row is None:
+ return {"success": False, "message": f"Episode {episode_index} not found in {repo_id}"}
+
+ cameras: list[dict[str, Any]] = []
+ for cam in episode_media.list_cameras(dataset_dir):
+ try:
+ loc = episode_media.locate_episode_video(dataset_dir, episode_index, camera=cam)
+ except episode_media.EpisodeNotFoundError as e:
+ logger.warning(f"{repo_id} ep{episode_index}: camera {cam} unavailable: {e}")
+ continue
+ cameras.append(
+ {
+ "name": cam,
+ "from_timestamp": loc.from_timestamp,
+ "to_timestamp": loc.to_timestamp,
+ "chunk": loc.chunk,
+ "file_index": loc.file_index,
+ }
+ )
+
+ indices = [r.episode_idx for r in rows]
+ position = indices.index(episode_index)
+
+ return {
+ "success": True,
+ "repo_id": repo_id,
+ "episode_index": episode_index,
+ "length": row.length,
+ "duration_s": row.duration_s,
+ "tasks": list(row.tasks),
+ "fps": info.get("fps"),
+ "robot_type": info.get("robot_type"),
+ "cameras": cameras,
+ "prev_episode": indices[position - 1] if position > 0 else None,
+ "next_episode": indices[position + 1] if position < len(indices) - 1 else None,
+ "total_episodes": len(indices),
+ }
+
+
+def load_episode_frame_png(
+ repo_id: str, episode_index: int, camera: str | None, frame_index: int, max_width: int | None
+) -> bytes:
+ """One decoded frame as PNG bytes. Raises on any miss; the route maps to 404."""
+ dataset_dir = episode_media.resolve_dataset_dir(repo_id)
+ location = episode_media.locate_episode_video(dataset_dir, episode_index, camera=camera)
+ return episode_media.extract_frame_png(location, frame_index, max_width=max_width)
+
+
+def handle_episode_thumbnails(
+ repo_id: str, episode_index: int, camera: str | None, count: int
+) -> dict[str, Any]:
+ """``count`` evenly-spaced thumbnails for the film strip, decoded in one pass."""
+ try:
+ dataset_dir = episode_media.resolve_dataset_dir(repo_id)
+ location = episode_media.locate_episode_video(dataset_dir, episode_index, camera=camera)
+ rows = episode_media.read_episode_index(dataset_dir)
+ except (episode_media.DatasetNotFoundError, episode_media.EpisodeNotFoundError) as e:
+ return {"success": False, "message": str(e)}
+
+ row = next((r for r in rows if r.episode_idx == episode_index), None)
+ length = row.length if row and row.length else 0
+ if length <= 0:
+ return {"success": False, "message": f"Episode {episode_index} has no frames"}
+
+ count = max(1, min(count, length))
+ # Evenly spaced across the whole episode, inclusive of both ends: a strip
+ # that stops short of the last frame hides how an episode finishes, which is
+ # usually the part you are checking.
+ last = length - 1
+ indices = [0] if count == 1 else sorted({round(i * last / (count - 1)) for i in range(count)})
+
+ try:
+ pngs = episode_media.extract_thumbnails(location, indices)
+ except episode_media.EpisodeNotFoundError as e:
+ return {"success": False, "message": str(e)}
+
+ return {
+ "success": True,
+ "episode_index": episode_index,
+ "camera": location.camera,
+ "thumbnails": [
+ {
+ "frame_index": idx,
+ "data_uri": "data:image/png;base64," + base64.b64encode(png).decode("ascii"),
+ }
+ for idx, png in zip(indices, pngs, strict=False)
+ ],
+ }
+
+
+def locate_video_file(repo_id: str, camera: str, chunk: int, file_index: int) -> Path:
+ """On-disk mp4 for one camera's chunk/file. Raises on any miss; the route 404s.
+
+ Addressed by *file* rather than by episode on purpose. Several episodes share
+ one mp4, so a file-addressed URL stays byte-identical as the viewer pages
+ between them — the browser reuses the already-buffered video and the player
+ only seeks. Keying on the episode instead would give each one its own URL for
+ the same bytes, forcing a reload and dropping playback on every switch.
+ """
+ dataset_dir = episode_media.resolve_dataset_dir(repo_id)
+ return episode_media.resolve_video_file(dataset_dir, camera, chunk, file_index)
+
+
+def handle_episode_motion(repo_id: str, episode_index: int) -> dict[str, Any]:
+ """Per-frame aggregate joint motion for one episode."""
+ try:
+ dataset_dir = episode_media.resolve_dataset_dir(repo_id)
+ trace = episode_media.episode_motion_trace(dataset_dir, episode_index)
+ except (episode_media.DatasetNotFoundError, episode_media.EpisodeNotFoundError) as e:
+ return {"success": False, "message": str(e)}
+ except Exception as e:
+ logger.warning(f"Could not read motion for {repo_id} ep{episode_index}: {e}")
+ return {"success": False, "message": f"Could not read action stream: {e}"}
+
+ return {
+ "success": True,
+ "episode_index": episode_index,
+ "motion": trace,
+ "max_motion": max(trace) if trace else 0.0,
+ }
diff --git a/lelab/episode_media.py b/lelab/episode_media.py
new file mode 100644
index 00000000..6e31c3f1
--- /dev/null
+++ b/lelab/episode_media.py
@@ -0,0 +1,518 @@
+# Copyright 2026 VibeCuisine. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Read-only episode media access for local LeRobot datasets.
+
+This is the read half of the dataset browser: it locates the mp4 holding a given
+episode/camera and decodes frames out of it. It deliberately reads the on-disk
+layout directly (``meta/episodes/*.parquet`` + ``videos/``) rather than going
+through ``LeRobotDataset``, so listing and previewing a dataset stays cheap and
+never constructs a robot config or touches torch.
+
+Decoding uses PyAV, which already ships with ``lerobot[dataset]``. Frames are
+decoded in-process; there is no ffmpeg subprocess.
+
+The layout this understands (LeRobot v3.0)::
+
+
/meta/info.json
+ /meta/episodes/chunk-000/file-000.parquet
+ /videos/observation.images./chunk-000/file-000.mp4
+
+Several episodes usually share one mp4. Each episode's slice of the file is
+recorded in the episodes parquet as ``.../from_timestamp`` and
+``.../to_timestamp``, so locating an episode means resolving both the file *and*
+the time window inside it. Older datasets store one mp4 per episode and carry no
+window; :func:`_locate_fallback` covers those.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import logging
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+VIDEO_KEY_PREFIX = "observation.images."
+
+
+class DatasetNotFoundError(LookupError):
+ """No local dataset directory for the requested repo_id."""
+
+
+class EpisodeNotFoundError(LookupError):
+ """The requested episode/camera/frame combination doesn't exist on disk."""
+
+
+# ── dataset resolution ──────────────────────────────────────────────────────
+
+
+def lerobot_cache_root() -> Path:
+ """Root of the local LeRobot dataset cache.
+
+ Reads the environment at call time rather than importing lerobot's
+ ``HF_LEROBOT_HOME`` constant, which is frozen at import. Same source of
+ truth as ``datasets._lerobot_cache_root``, and it keeps this module free of
+ a torch-pulling import it doesn't otherwise need.
+ """
+ return Path(os.environ.get("HF_LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser().resolve()
+
+
+def resolve_dataset_dir(repo_id: str) -> Path:
+ """Resolve ``repo_id`` to a local dataset directory.
+
+ Rejects path traversal the same way ``handle_delete_dataset`` does: the
+ resolved target must stay strictly inside the cache root. ``repo_id`` comes
+ from a query string, so this is load-bearing rather than defensive.
+ """
+ root = lerobot_cache_root()
+ target = (root / repo_id).resolve()
+ if target == root or root not in target.parents:
+ raise DatasetNotFoundError(f"Invalid dataset path: {repo_id}")
+ if not (target / "meta" / "info.json").is_file():
+ raise DatasetNotFoundError(f"Dataset not found locally: {repo_id}")
+ return target
+
+
+def read_info(dataset_dir: Path) -> dict[str, Any]:
+ """Parsed ``meta/info.json``."""
+ try:
+ with open(dataset_dir / "meta" / "info.json") as f:
+ return json.load(f)
+ except (OSError, json.JSONDecodeError) as e:
+ raise DatasetNotFoundError(f"Unreadable meta/info.json in {dataset_dir}: {e}") from e
+
+
+def list_cameras(dataset_dir: Path) -> list[str]:
+ """Camera names stored as video features (``observation.images.``)."""
+ features = read_info(dataset_dir).get("features") or {}
+ return sorted(
+ key.removeprefix(VIDEO_KEY_PREFIX)
+ for key, feat in features.items()
+ if key.startswith(VIDEO_KEY_PREFIX) and isinstance(feat, dict) and feat.get("dtype") == "video"
+ )
+
+
+# ── episode index ───────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class EpisodeIndexRow:
+ episode_idx: int
+ length: int
+ duration_s: float | None
+ from_timestamp: float | None
+ to_timestamp: float | None
+ tasks: tuple[str, ...]
+
+
+def _maybe_float(value: Any) -> float | None:
+ if value is None:
+ return None
+ try:
+ return float(value)
+ except (TypeError, ValueError):
+ return None
+
+
+def read_episode_index(dataset_dir: Path) -> list[EpisodeIndexRow]:
+ """Per-episode metadata from ``meta/episodes/*.parquet``, sorted by index.
+
+ ``duration_s`` is derived from ``length / fps``, *not* from the video's
+ ``to_timestamp - from_timestamp`` window. An episode the operator accepted
+ early stops its parquet at accept while its video still spans the full
+ configured window, so the window over-reports the demonstration — a 2.6s
+ episode reads as the configured 10.0s. The window is kept on
+ ``from_/to_timestamp`` for the video transport and used for ``duration_s``
+ only when a row has no usable length/fps.
+
+ Returns ``[]`` rather than raising when the meta dir is missing or
+ unreadable, so a partially-synced dataset still lists.
+ """
+ import pyarrow.parquet as pq
+
+ eps_dir = dataset_dir / "meta" / "episodes"
+ if not eps_dir.is_dir():
+ return []
+
+ info = read_info(dataset_dir)
+ fps = float(info.get("fps") or 0.0)
+ cameras = list_cameras(dataset_dir)
+ # The window is read from the first camera so the listing is stable
+ # regardless of camera; in practice every camera shares the same window.
+ cam = cameras[0] if cameras else None
+ from_col = f"videos/{VIDEO_KEY_PREFIX}{cam}/from_timestamp" if cam else None
+ to_col = f"videos/{VIDEO_KEY_PREFIX}{cam}/to_timestamp" if cam else None
+
+ rows: list[EpisodeIndexRow] = []
+ for f in sorted(eps_dir.rglob("*.parquet")):
+ try:
+ schema = pq.read_schema(f)
+ except Exception as e:
+ logger.warning(f"Skipping unreadable episodes parquet {f}: {e}")
+ continue
+ wanted = [
+ c for c in ("episode_index", "length", "tasks", from_col, to_col) if c and c in schema.names
+ ]
+ if "episode_index" not in wanted:
+ continue
+ try:
+ t = pq.read_table(f, columns=wanted).to_pydict()
+ except Exception as e:
+ logger.warning(f"Skipping unreadable episodes parquet {f}: {e}")
+ continue
+
+ n = len(t["episode_index"])
+ for i in range(n):
+ length = int(t.get("length", [0] * n)[i] or 0)
+ ft = _maybe_float(t[from_col][i]) if from_col in t else None
+ tt = _maybe_float(t[to_col][i]) if to_col in t else None
+ if length and fps:
+ duration = length / fps
+ elif ft is not None and tt is not None:
+ duration = max(0.0, tt - ft)
+ else:
+ duration = None
+ tasks_raw = t.get("tasks", [[]] * n)[i] if "tasks" in t else []
+ rows.append(
+ EpisodeIndexRow(
+ episode_idx=int(t["episode_index"][i]),
+ length=length,
+ duration_s=duration,
+ from_timestamp=ft,
+ to_timestamp=tt,
+ tasks=tuple(str(x) for x in (tasks_raw or [])),
+ )
+ )
+
+ rows.sort(key=lambda r: r.episode_idx)
+ return rows
+
+
+# ── video location ──────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class EpisodeVideoLocation:
+ camera: str
+ chunk: int
+ file_index: int
+ episode_idx: int
+ path: Path
+ fps: float
+ # None when one episode owns the whole file (older layouts): the episode
+ # spans the entire video rather than a window inside it.
+ from_timestamp: float | None = None
+ to_timestamp: float | None = None
+
+
+def _locate_via_meta(
+ dataset_dir: Path, episode_idx: int, camera: str, fps: float
+) -> EpisodeVideoLocation | None:
+ import pyarrow.parquet as pq
+
+ eps_dir = dataset_dir / "meta" / "episodes"
+ if not eps_dir.is_dir():
+ return None
+
+ base = f"videos/{VIDEO_KEY_PREFIX}{camera}"
+ chunk_col, file_col = f"{base}/chunk_index", f"{base}/file_index"
+ from_col, to_col = f"{base}/from_timestamp", f"{base}/to_timestamp"
+
+ for f in sorted(eps_dir.rglob("*.parquet")):
+ try:
+ schema = pq.read_schema(f)
+ except Exception:
+ continue
+ wanted = ["episode_index"] + [c for c in (chunk_col, file_col, from_col, to_col) if c in schema.names]
+ if "episode_index" not in schema.names:
+ continue
+ try:
+ t = pq.read_table(f, columns=wanted).to_pydict()
+ except Exception:
+ continue
+
+ for i, ep in enumerate(t.get("episode_index", [])):
+ if int(ep) != episode_idx:
+ continue
+ chunk = int(t[chunk_col][i]) if chunk_col in t else 0
+ file_idx = int(t[file_col][i]) if file_col in t else episode_idx
+ return EpisodeVideoLocation(
+ camera=camera,
+ chunk=chunk,
+ file_index=file_idx,
+ episode_idx=episode_idx,
+ path=(
+ dataset_dir
+ / "videos"
+ / f"{VIDEO_KEY_PREFIX}{camera}"
+ / f"chunk-{chunk:03d}"
+ / f"file-{file_idx:03d}.mp4"
+ ),
+ fps=fps,
+ from_timestamp=_maybe_float(t[from_col][i]) if from_col in t else None,
+ to_timestamp=_maybe_float(t[to_col][i]) if to_col in t else None,
+ )
+ return None
+
+
+def _resolve_action_file(dataset_dir: Path, episode_idx: int) -> Path | None:
+ """Which ``data/chunk-N/file-M.parquet`` holds this episode's frames.
+
+ The episodes meta pins the chunk/file for the action stream the same way it
+ does for video. Older datasets don't, so fall back to the single-file layout.
+ """
+ import pyarrow.parquet as pq
+
+ eps_dir = dataset_dir / "meta" / "episodes"
+ if eps_dir.is_dir():
+ for f in sorted(eps_dir.rglob("*.parquet")):
+ try:
+ schema = pq.read_schema(f)
+ except Exception:
+ continue
+ if "episode_index" not in schema.names:
+ continue
+ wanted = ["episode_index"] + [
+ c for c in ("data/chunk_index", "data/file_index") if c in schema.names
+ ]
+ try:
+ t = pq.read_table(f, columns=wanted).to_pydict()
+ except Exception:
+ continue
+ for i, ep in enumerate(t.get("episode_index", [])):
+ if int(ep) != episode_idx:
+ continue
+ chunk = int(t["data/chunk_index"][i]) if "data/chunk_index" in t else 0
+ file_idx = int(t["data/file_index"][i]) if "data/file_index" in t else 0
+ target = dataset_dir / "data" / f"chunk-{chunk:03d}" / f"file-{file_idx:03d}.parquet"
+ if target.is_file():
+ return target
+
+ fallback = dataset_dir / "data" / "chunk-000" / "file-000.parquet"
+ return fallback if fallback.is_file() else None
+
+
+def load_episode_actions(dataset_dir: Path, episode_idx: int) -> list[list[float]]:
+ """This episode's per-frame commanded action vectors, in file order.
+
+ Returns a ``(T, dof)`` nested list. Kept as plain Python rather than numpy so
+ the caller can hand it straight to the JSON encoder.
+ """
+ import pyarrow.parquet as pq
+
+ file_path = _resolve_action_file(dataset_dir, episode_idx)
+ if file_path is None:
+ raise EpisodeNotFoundError(f"No data parquet for episode {episode_idx}")
+ try:
+ table = pq.read_table(file_path, columns=["episode_index", "action"])
+ except Exception as e:
+ raise EpisodeNotFoundError(f"Could not read {file_path.name}: {e}") from e
+
+ if "action" not in table.column_names:
+ raise EpisodeNotFoundError(f"{file_path.name} has no 'action' column")
+
+ indices = table["episode_index"].to_pylist()
+ actions = table["action"].to_pylist()
+ rows = [a for ep, a in zip(indices, actions, strict=False) if int(ep) == episode_idx]
+ if not rows:
+ raise EpisodeNotFoundError(f"Episode {episode_idx} has no action rows")
+ # Some datasets store a scalar action per frame; normalise to 2-D.
+ return [list(r) if isinstance(r, (list, tuple)) else [r] for r in rows]
+
+
+def episode_motion_trace(dataset_dir: Path, episode_idx: int) -> list[float]:
+ """Per-frame aggregate joint motion: the sum of absolute joint deltas.
+
+ One scalar per frame collapsing the whole action vector, so a single trace
+ shows where the arm was busy and where it sat idle — which is what makes the
+ idle head and tail of a demonstration visible at a glance. Frame 0 has no
+ predecessor and is reported as 0.
+ """
+ actions = load_episode_actions(dataset_dir, episode_idx)
+ trace = [0.0]
+ for prev, cur in zip(actions, actions[1:], strict=False):
+ trace.append(float(sum(abs(c - p) for p, c in zip(prev, cur, strict=False))))
+ return trace
+
+
+def _locate_fallback(
+ dataset_dir: Path, episode_idx: int, camera: str, fps: float
+) -> EpisodeVideoLocation | None:
+ """Pre-packing layouts: one mp4 per episode, no time window."""
+ cam_dir = dataset_dir / "videos" / f"{VIDEO_KEY_PREFIX}{camera}"
+ for candidate in (
+ cam_dir / "chunk-000" / f"file-{episode_idx:03d}.mp4",
+ cam_dir / f"episode_{episode_idx:06d}.mp4",
+ ):
+ if candidate.exists():
+ return EpisodeVideoLocation(
+ camera=camera,
+ chunk=0,
+ file_index=episode_idx,
+ episode_idx=episode_idx,
+ path=candidate,
+ fps=fps,
+ )
+ return None
+
+
+def resolve_video_file(dataset_dir: Path, camera: str, chunk: int, file_index: int) -> Path:
+ """The mp4 at ``videos/observation.images./chunk-N/file-M.mp4``.
+
+ ``chunk``/``file_index`` are ints, so the path can't escape the dataset; the
+ camera is checked against the declared video features rather than
+ interpolated blind.
+ """
+ if camera not in list_cameras(dataset_dir):
+ raise EpisodeNotFoundError(f"Camera {camera!r} not in this dataset")
+ path = (
+ dataset_dir
+ / "videos"
+ / f"{VIDEO_KEY_PREFIX}{camera}"
+ / f"chunk-{int(chunk):03d}"
+ / f"file-{int(file_index):03d}.mp4"
+ )
+ if not path.is_file():
+ raise EpisodeNotFoundError(f"No video at {path.name} for camera {camera!r}")
+ return path
+
+
+def locate_episode_video(
+ dataset_dir: Path, episode_idx: int, *, camera: str | None = None
+) -> EpisodeVideoLocation:
+ """Resolve the mp4 (and time window) holding ``episode_idx`` for one camera."""
+ info = read_info(dataset_dir)
+ fps = float(info.get("fps") or 0.0)
+ cameras = list_cameras(dataset_dir)
+ if not cameras:
+ raise EpisodeNotFoundError(f"No camera videos in {dataset_dir.name}")
+ cam = camera or cameras[0]
+ if cam not in cameras:
+ raise EpisodeNotFoundError(f"Camera {cam!r} not in {cameras}")
+
+ location = _locate_via_meta(dataset_dir, episode_idx, cam, fps) or _locate_fallback(
+ dataset_dir, episode_idx, cam, fps
+ )
+ if location is None or not location.path.exists():
+ raise EpisodeNotFoundError(f"Episode {episode_idx} for camera {cam!r} not found on disk")
+ return location
+
+
+# ── decoding ────────────────────────────────────────────────────────────────
+
+
+def _frame_time(location: EpisodeVideoLocation, frame_index: int) -> float:
+ """Absolute timestamp *inside the mp4* for an episode-relative frame index."""
+ offset = location.from_timestamp or 0.0
+ if not location.fps:
+ return offset
+ return offset + (frame_index / location.fps)
+
+
+def _decode_at(container, stream, target_s: float):
+ """Decode the first frame at or after ``target_s``.
+
+ Seeks to the keyframe at or before the target, then walks forward. Returns
+ the last decoded frame if the target is past the end (a short episode
+ shouldn't 500 on its final thumbnail).
+ """
+ # Every real mp4 stream carries a time_base; the fallback keeps a synthetic
+ # or damaged stream from dividing by None.
+ seek_target = int(target_s / stream.time_base) if stream.time_base else int(target_s * 1_000_000)
+ try:
+ container.seek(max(seek_target, 0), stream=stream, backward=True, any_frame=False)
+ except Exception:
+ container.seek(0)
+
+ last = None
+ for frame in container.decode(stream):
+ last = frame
+ if frame.time is not None and frame.time >= target_s - 1e-4:
+ return frame
+ return last
+
+
+def extract_frame_png(
+ location: EpisodeVideoLocation, frame_index: int, *, max_width: int | None = None
+) -> bytes:
+ """Decode one episode-relative frame to PNG bytes."""
+ import av
+ from PIL import Image
+
+ target = _frame_time(location, frame_index)
+ try:
+ with av.open(str(location.path)) as container:
+ stream = container.streams.video[0]
+ stream.thread_type = "AUTO"
+ frame = _decode_at(container, stream, target)
+ if frame is None:
+ raise EpisodeNotFoundError(f"No frame at {target:.3f}s in {location.path.name}")
+ img = frame.to_image()
+ except EpisodeNotFoundError:
+ raise
+ except Exception as e:
+ raise EpisodeNotFoundError(f"Could not decode {location.path.name}: {e}") from e
+
+ if max_width and img.width > max_width:
+ height = max(1, round(img.height * max_width / img.width))
+ img = img.resize((max_width, height), Image.BILINEAR)
+
+ buf = io.BytesIO()
+ img.save(buf, format="PNG")
+ return buf.getvalue()
+
+
+def extract_thumbnails(
+ location: EpisodeVideoLocation,
+ frame_indices: list[int],
+ *,
+ max_width: int = 160,
+) -> list[bytes]:
+ """Decode several frames in a single pass over the file.
+
+ The film strip wants N frames from one episode; opening, seeking and tearing
+ down the container N times is the obvious way and the wasteful one. Frames
+ are requested in ascending order, so one forward walk serves them all.
+ """
+ import av
+ from PIL import Image
+
+ targets = sorted({int(i) for i in frame_indices})
+ if not targets:
+ return []
+
+ out: list[bytes] = []
+ try:
+ with av.open(str(location.path)) as container:
+ stream = container.streams.video[0]
+ stream.thread_type = "AUTO"
+ for idx in targets:
+ frame = _decode_at(container, stream, _frame_time(location, idx))
+ if frame is None:
+ continue
+ img = frame.to_image()
+ if img.width > max_width:
+ height = max(1, round(img.height * max_width / img.width))
+ img = img.resize((max_width, height), Image.BILINEAR)
+ buf = io.BytesIO()
+ img.save(buf, format="PNG")
+ out.append(buf.getvalue())
+ except Exception as e:
+ raise EpisodeNotFoundError(f"Could not decode {location.path.name}: {e}") from e
+ return out
diff --git a/lelab/server.py b/lelab/server.py
index 519e0c72..279f419b 100644
--- a/lelab/server.py
+++ b/lelab/server.py
@@ -28,7 +28,7 @@
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from starlette.datastructures import Headers
@@ -40,6 +40,7 @@
# Import our custom calibration functionality
from .calibrate import CalibrationRequest, calibration_manager
+from .episode_media import DatasetNotFoundError, EpisodeNotFoundError
from .jobs import (
JobAlreadyRunningError,
JobNotFoundError,
@@ -387,6 +388,76 @@ def datasets_list():
return dataset_browser.list_all_datasets()
+@app.get("/dataset-episodes")
+def dataset_episodes(repo_id: str):
+ """List the episodes of a local dataset."""
+ return dataset_browser.handle_list_episodes(repo_id)
+
+
+@app.get("/dataset-episode")
+def dataset_episode(repo_id: str, episode_index: int):
+ """One episode's metadata and per-camera video windows."""
+ return dataset_browser.handle_episode_detail(repo_id, episode_index)
+
+
+@app.get("/dataset-motion")
+def dataset_motion(repo_id: str, episode_index: int):
+ """Per-frame aggregate joint motion (sum of absolute joint deltas)."""
+ return dataset_browser.handle_episode_motion(repo_id, episode_index)
+
+
+@app.get("/dataset-thumbnails")
+def dataset_thumbnails(repo_id: str, episode_index: int, camera: str | None = None, count: int = 12):
+ """Evenly-spaced film-strip thumbnails, decoded in a single pass over the mp4."""
+ return dataset_browser.handle_episode_thumbnails(repo_id, episode_index, camera, count)
+
+
+@app.get("/dataset-frame")
+def dataset_frame(
+ repo_id: str,
+ episode_index: int,
+ camera: str | None = None,
+ frame_index: int = 0,
+ max_width: int | None = None,
+):
+ """A single decoded frame as a PNG."""
+ try:
+ png = dataset_browser.load_episode_frame_png(repo_id, episode_index, camera, frame_index, max_width)
+ except (DatasetNotFoundError, EpisodeNotFoundError) as exc:
+ raise HTTPException(status_code=404, detail=str(exc)) from exc
+ # A frame is immutable for a given (dataset, episode, camera, index), and the
+ # scrubber re-requests the same ones constantly, so let the browser cache it.
+ return Response(
+ content=png,
+ media_type="image/png",
+ headers={"Cache-Control": "private, max-age=3600"},
+ )
+
+
+@app.get("/dataset-video")
+def dataset_video(repo_id: str, camera: str, chunk: int = 0, file: int = 0):
+ """One camera's mp4, served whole for in-page playback.
+
+ Addressed by chunk/file rather than by episode: several episodes share a
+ file, so this URL stays stable as the viewer pages between them and the
+ browser reuses what it already buffered. Callers get the episode's window
+ from /dataset-episode and seek to it. FileResponse answers Range requests,
+ so the browser does its own seeking rather than us transcoding a cut.
+ """
+ try:
+ path = dataset_browser.locate_video_file(repo_id, camera, chunk, file)
+ except (DatasetNotFoundError, EpisodeNotFoundError) as exc:
+ raise HTTPException(status_code=404, detail=str(exc)) from exc
+ # content_disposition_type="inline" is load-bearing: the default is
+ # "attachment", which makes the browser download the file rather than play it.
+ return FileResponse(
+ path,
+ media_type="video/mp4",
+ filename=path.name,
+ content_disposition_type="inline",
+ )
+
+
@app.get("/ws-test")
def websocket_test():
"""Test endpoint to verify WebSocket support"""
@@ -784,7 +855,9 @@ def run_update():
return handle_run_update()
-# Replay is rendered by the embedded lerobot/visualize_dataset Space; no backend routes needed.
+# Datasets with files on disk are browsed locally (see the /dataset-* routes above).
+# Hub-only datasets have nothing here to decode, so the frontend still hands those
+# off to the lerobot/visualize_dataset Space.
# ============================================================================
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 2b7cf7f0..697c817f 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -11,13 +11,25 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Tests for lelab.datasets — local cache walk and merge logic."""
+"""Tests for lelab.datasets — local cache walk, merge logic, episode browsing."""
from __future__ import annotations
from pathlib import Path
from unittest.mock import patch
+import pytest
+from fastapi.testclient import TestClient
+
+from .test_episode_media import _write_dataset
+
+
+@pytest.fixture
+def browsable_dataset(tmp_lerobot_home: Path) -> str:
+ """A real v3.0 dataset on disk: 3 episodes, 2 cameras, 10fps."""
+ _write_dataset(tmp_lerobot_home, "acme/demo", episodes=[(0, 12), (1, 8), (2, 15)])
+ return "acme/demo"
+
def _make_dataset(root: Path, repo_id: str) -> None:
"""Create the minimal layout `_is_dataset_dir` recognizes."""
@@ -105,3 +117,149 @@ def test_list_all_datasets_merges_hub_and_local(
by_id = {d["repo_id"]: d for d in result}
assert by_id["alice/pusht"]["source"] == "both"
assert by_id["alice/aloha"]["source"] == "hub"
+
+
+# ── episode browsing ────────────────────────────────────────────────────────
+#
+# The handlers are covered end-to-end through the routes, so the query-param
+# contract and the 404 mapping for binary endpoints are checked too. The
+# dataset fixture builds a real v3.0 layout (see tests/test_episode_media.py).
+
+
+def test_list_episodes_unknown_dataset(client: TestClient, tmp_lerobot_home: Path) -> None:
+ r = client.get("/dataset-episodes", params={"repo_id": "nope/nope"})
+ assert r.status_code == 200
+ assert r.json()["success"] is False
+
+
+def test_list_episodes_returns_rows(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get("/dataset-episodes", params={"repo_id": browsable_dataset})
+ body = r.json()
+ assert body["success"] is True
+ assert body["fps"] == 10
+ assert body["cameras"] == ["top", "wrist"]
+ assert [e["episode_index"] for e in body["episodes"]] == [0, 1, 2]
+
+
+def test_episode_detail_reports_neighbours(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get("/dataset-episode", params={"repo_id": browsable_dataset, "episode_index": 1})
+ body = r.json()
+ assert body["success"] is True
+ assert body["prev_episode"] == 0
+ assert body["next_episode"] == 2
+ assert {c["name"] for c in body["cameras"]} == {"top", "wrist"}
+
+
+def test_episode_detail_unknown_episode(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get("/dataset-episode", params={"repo_id": browsable_dataset, "episode_index": 99})
+ assert r.status_code == 200
+ assert r.json()["success"] is False
+
+
+def test_frame_route_returns_png(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get(
+ "/dataset-frame",
+ params={"repo_id": browsable_dataset, "episode_index": 0, "camera": "top", "frame_index": 2},
+ )
+ assert r.status_code == 200
+ assert r.headers["content-type"] == "image/png"
+ assert r.content[:4] == b"\x89PNG"
+
+
+def test_frame_route_404s_on_traversal(client: TestClient, tmp_lerobot_home: Path) -> None:
+ r = client.get("/dataset-frame", params={"repo_id": "../../etc", "episode_index": 0})
+ assert r.status_code == 404
+
+
+def test_frame_route_404s_on_unknown_episode(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get("/dataset-frame", params={"repo_id": browsable_dataset, "episode_index": 99})
+ assert r.status_code == 404
+
+
+def test_video_route_is_inline_and_range_capable(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get(
+ "/dataset-video", params={"repo_id": browsable_dataset, "camera": "top", "chunk": 0, "file": 0}
+ )
+ assert r.status_code == 200
+ assert r.headers["content-type"] == "video/mp4"
+ # "attachment" would make the browser download the file instead of playing it.
+ assert "inline" in r.headers["content-disposition"]
+ assert r.headers.get("accept-ranges") == "bytes"
+
+
+def test_video_route_serves_a_range(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get(
+ "/dataset-video",
+ params={"repo_id": browsable_dataset, "camera": "top", "chunk": 0, "file": 0},
+ headers={"Range": "bytes=0-31"},
+ )
+ assert r.status_code == 206
+ assert len(r.content) == 32
+
+
+def test_video_url_is_shared_by_episodes_in_one_file(client: TestClient, browsable_dataset: str) -> None:
+ """Every episode of this dataset lives in one mp4, so they must share a URL.
+
+ The video route is addressed by chunk/file precisely so the player can page
+ between episodes without the src changing — an episode-addressed URL would
+ hand the browser a new identifier for the same bytes, forcing a reload and
+ dropping playback on every switch.
+ """
+ detail = [
+ client.get("/dataset-episode", params={"repo_id": browsable_dataset, "episode_index": i}).json()
+ for i in (0, 1, 2)
+ ]
+ top = [next(c for c in d["cameras"] if c["name"] == "top") for d in detail]
+ assert {(c["chunk"], c["file_index"]) for c in top} == {(0, 0)}
+ # ...while each still reports its own window inside that shared file.
+ assert len({c["from_timestamp"] for c in top}) == 3
+
+
+def test_video_route_404s_on_unknown_camera(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get(
+ "/dataset-video", params={"repo_id": browsable_dataset, "camera": "nope", "chunk": 0, "file": 0}
+ )
+ assert r.status_code == 404
+
+
+def test_video_route_404s_on_missing_file(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get(
+ "/dataset-video", params={"repo_id": browsable_dataset, "camera": "top", "chunk": 9, "file": 9}
+ )
+ assert r.status_code == 404
+
+
+def test_thumbnails_span_the_whole_episode(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get(
+ "/dataset-thumbnails",
+ params={"repo_id": browsable_dataset, "episode_index": 0, "camera": "top", "count": 4},
+ )
+ body = r.json()
+ assert body["success"] is True
+ frames = [t["frame_index"] for t in body["thumbnails"]]
+ # Episode 0 is 12 frames: the strip must reach the last one, not stop short.
+ assert frames[0] == 0
+ assert frames[-1] == 11
+ assert all(t["data_uri"].startswith("data:image/png;base64,") for t in body["thumbnails"])
+
+
+def test_thumbnails_clamped_to_episode_length(client: TestClient, browsable_dataset: str) -> None:
+ # Episode 1 has 8 frames; asking for 50 must not invent any.
+ r = client.get(
+ "/dataset-thumbnails",
+ params={"repo_id": browsable_dataset, "episode_index": 1, "count": 50},
+ )
+ assert len(r.json()["thumbnails"]) <= 8
+
+
+def test_motion_route_returns_one_value_per_frame(client: TestClient, browsable_dataset: str) -> None:
+ r = client.get("/dataset-motion", params={"repo_id": browsable_dataset, "episode_index": 0})
+ body = r.json()
+ assert body["success"] is True
+ assert len(body["motion"]) == 12
+ assert body["motion"][0] == 0.0
+
+
+def test_motion_route_unknown_dataset(client: TestClient, tmp_lerobot_home: Path) -> None:
+ r = client.get("/dataset-motion", params={"repo_id": "nope/nope", "episode_index": 0})
+ assert r.json()["success"] is False
diff --git a/tests/test_episode_media.py b/tests/test_episode_media.py
new file mode 100644
index 00000000..cf62b7fd
--- /dev/null
+++ b/tests/test_episode_media.py
@@ -0,0 +1,338 @@
+# Copyright 2026 VibeCuisine. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for lelab.episode_media — on-disk layout reads and frame decoding.
+
+Builds a real (tiny) LeRobot v3.0 dataset on disk — a genuine mp4 encoded with
+PyAV plus the parquet metadata — rather than mocking the readers. The whole
+point of this module is that it agrees with the on-disk format, which a mock
+cannot check.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+
+
+def _write_dataset(
+ root: Path,
+ repo_id: str,
+ *,
+ episodes: list[tuple[int, int]], # (episode_index, length)
+ cameras: tuple[str, ...] = ("top", "wrist"),
+ fps: int = 10,
+ dof: int = 3,
+ packed: bool = True,
+) -> Path:
+ """Write a minimal but real v3.0 dataset. Returns its directory.
+
+ ``packed=True`` puts every episode in one mp4 with from/to windows (the
+ modern layout); ``packed=False`` omits the windows so the fallback path is
+ exercised.
+ """
+ import av
+ import numpy as np
+ import pyarrow as pa
+ import pyarrow.parquet as pq
+
+ d = root / repo_id
+ (d / "meta" / "episodes" / "chunk-000").mkdir(parents=True)
+ (d / "data" / "chunk-000").mkdir(parents=True)
+
+ features = {f"observation.images.{c}": {"dtype": "video", "shape": [24, 32, 3]} for c in cameras}
+ features["action"] = {"dtype": "float32", "shape": [dof]}
+ (d / "meta" / "info.json").write_text(
+ __import__("json").dumps(
+ {
+ "codebase_version": "v3.0",
+ "fps": fps,
+ "robot_type": "test_follower",
+ "total_episodes": len(episodes),
+ "total_frames": sum(n for _, n in episodes),
+ "features": features,
+ }
+ )
+ )
+
+ total = sum(n for _, n in episodes)
+ for cam in cameras:
+ vdir = d / "videos" / f"observation.images.{cam}" / "chunk-000"
+ vdir.mkdir(parents=True)
+ with av.open(str(vdir / "file-000.mp4"), "w") as c:
+ s = c.add_stream("libx264", rate=fps)
+ s.width, s.height, s.pix_fmt = 32, 24, "yuv420p"
+ for i in range(total):
+ # A per-frame ramp so a decoded frame is identifiable.
+ img = np.full((24, 32, 3), (i * 3) % 256, dtype=np.uint8)
+ c.mux(s.encode(av.VideoFrame.from_ndarray(img, format="rgb24")))
+ c.mux(s.encode(None))
+
+ # Episode metadata: windows laid end to end, matching how episodes are packed.
+ meta: dict[str, list] = {"episode_index": [], "length": [], "tasks": []}
+ for cam in cameras:
+ meta[f"videos/observation.images.{cam}/chunk_index"] = []
+ meta[f"videos/observation.images.{cam}/file_index"] = []
+ if packed:
+ meta[f"videos/observation.images.{cam}/from_timestamp"] = []
+ meta[f"videos/observation.images.{cam}/to_timestamp"] = []
+ meta["data/chunk_index"] = []
+ meta["data/file_index"] = []
+
+ cursor = 0.0
+ for ep, n in episodes:
+ meta["episode_index"].append(ep)
+ meta["length"].append(n)
+ meta["tasks"].append(["do the thing"])
+ meta["data/chunk_index"].append(0)
+ meta["data/file_index"].append(0)
+ for cam in cameras:
+ meta[f"videos/observation.images.{cam}/chunk_index"].append(0)
+ meta[f"videos/observation.images.{cam}/file_index"].append(0)
+ if packed:
+ meta[f"videos/observation.images.{cam}/from_timestamp"].append(cursor)
+ meta[f"videos/observation.images.{cam}/to_timestamp"].append(cursor + n / fps)
+ cursor += n / fps
+ pq.write_table(pa.table(meta), d / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
+
+ # Frames: a per-joint ramp so motion is a known constant.
+ ep_col, act_col = [], []
+ for ep, n in episodes:
+ for i in range(n):
+ ep_col.append(ep)
+ act_col.append([float(i)] * dof)
+ pq.write_table(
+ pa.table({"episode_index": ep_col, "action": act_col}),
+ d / "data" / "chunk-000" / "file-000.parquet",
+ )
+ return d
+
+
+@pytest.fixture
+def dataset(tmp_lerobot_home: Path) -> str:
+ _write_dataset(tmp_lerobot_home, "acme/demo", episodes=[(0, 12), (1, 8), (2, 15)])
+ return "acme/demo"
+
+
+# ── dataset resolution ──────────────────────────────────────────────────────
+
+
+def test_resolve_dataset_dir_rejects_traversal(tmp_lerobot_home: Path) -> None:
+ from lelab.episode_media import DatasetNotFoundError, resolve_dataset_dir
+
+ with pytest.raises(DatasetNotFoundError):
+ resolve_dataset_dir("../../../etc")
+
+
+def test_resolve_dataset_dir_rejects_root_itself(tmp_lerobot_home: Path) -> None:
+ from lelab.episode_media import DatasetNotFoundError, resolve_dataset_dir
+
+ with pytest.raises(DatasetNotFoundError):
+ resolve_dataset_dir(".")
+
+
+def test_resolve_dataset_dir_rejects_non_dataset(tmp_lerobot_home: Path) -> None:
+ from lelab.episode_media import DatasetNotFoundError, resolve_dataset_dir
+
+ (tmp_lerobot_home / "not-a-dataset").mkdir()
+ with pytest.raises(DatasetNotFoundError):
+ resolve_dataset_dir("not-a-dataset")
+
+
+def test_list_cameras_only_returns_video_features(dataset: str) -> None:
+ from lelab.episode_media import list_cameras, resolve_dataset_dir
+
+ # `action` is a feature too, but it isn't a video and must not appear.
+ assert list_cameras(resolve_dataset_dir(dataset)) == ["top", "wrist"]
+
+
+# ── episode index ───────────────────────────────────────────────────────────
+
+
+def test_read_episode_index_duration_is_length_over_fps(dataset: str) -> None:
+ from lelab.episode_media import read_episode_index, resolve_dataset_dir
+
+ rows = read_episode_index(resolve_dataset_dir(dataset))
+ assert [r.episode_idx for r in rows] == [0, 1, 2]
+ assert [r.length for r in rows] == [12, 8, 15]
+ assert [r.duration_s for r in rows] == [1.2, 0.8, 1.5]
+ assert rows[0].tasks == ("do the thing",)
+
+
+def test_read_episode_index_prefers_length_over_video_window(tmp_lerobot_home: Path) -> None:
+ """An early-accepted episode's window over-reports the demonstration.
+
+ The window says 5s because that's how long the camera ran; the parquet says
+ 10 frames at 10fps = 1s, which is the honest duration. Length must win.
+ """
+ import pyarrow as pa
+ import pyarrow.parquet as pq
+
+ from lelab.episode_media import read_episode_index, resolve_dataset_dir
+
+ d = _write_dataset(tmp_lerobot_home, "acme/early", episodes=[(0, 10)], cameras=("top",))
+ meta_path = d / "meta" / "episodes" / "chunk-000" / "file-000.parquet"
+ t = pq.read_table(meta_path).to_pydict()
+ t["videos/observation.images.top/to_timestamp"] = [5.0] # camera ran 5s
+ pq.write_table(pa.table(t), meta_path)
+
+ row = read_episode_index(resolve_dataset_dir("acme/early"))[0]
+ assert row.duration_s == 1.0 # not 5.0
+ assert row.to_timestamp == 5.0 # window still exposed for the transport
+
+
+def test_read_episode_index_empty_when_meta_missing(tmp_lerobot_home: Path) -> None:
+ from lelab.episode_media import read_episode_index
+
+ d = tmp_lerobot_home / "bare"
+ (d / "meta").mkdir(parents=True)
+ (d / "meta" / "info.json").write_text("{}")
+ assert read_episode_index(d) == []
+
+
+# ── video location ──────────────────────────────────────────────────────────
+
+
+def test_locate_episode_video_resolves_window(dataset: str) -> None:
+ from lelab.episode_media import locate_episode_video, resolve_dataset_dir
+
+ d = resolve_dataset_dir(dataset)
+ # Episode 1 follows episode 0 (12 frames @10fps = 1.2s) in the same file.
+ loc = locate_episode_video(d, 1, camera="top")
+ assert loc.path.name == "file-000.mp4"
+ assert loc.from_timestamp == pytest.approx(1.2)
+ assert loc.to_timestamp == pytest.approx(2.0)
+ assert loc.fps == 10
+
+
+def test_locate_episode_video_unknown_camera(dataset: str) -> None:
+ from lelab.episode_media import EpisodeNotFoundError, locate_episode_video, resolve_dataset_dir
+
+ with pytest.raises(EpisodeNotFoundError):
+ locate_episode_video(resolve_dataset_dir(dataset), 0, camera="nope")
+
+
+def test_locate_episode_video_missing_episode(dataset: str) -> None:
+ from lelab.episode_media import EpisodeNotFoundError, locate_episode_video, resolve_dataset_dir
+
+ with pytest.raises(EpisodeNotFoundError):
+ locate_episode_video(resolve_dataset_dir(dataset), 99)
+
+
+def test_locate_episode_video_defaults_to_first_camera(dataset: str) -> None:
+ from lelab.episode_media import locate_episode_video, resolve_dataset_dir
+
+ assert locate_episode_video(resolve_dataset_dir(dataset), 0).camera == "top"
+
+
+def test_locate_falls_back_to_unpacked_layout(tmp_lerobot_home: Path) -> None:
+ """Older datasets carry no from/to window: the episode is the whole file."""
+ from lelab.episode_media import locate_episode_video, resolve_dataset_dir
+
+ _write_dataset(tmp_lerobot_home, "acme/old", episodes=[(0, 6)], cameras=("top",), packed=False)
+ loc = locate_episode_video(resolve_dataset_dir("acme/old"), 0, camera="top")
+ assert loc.from_timestamp is None
+ assert loc.to_timestamp is None
+ assert loc.path.exists()
+
+
+# ── decoding ────────────────────────────────────────────────────────────────
+
+
+def test_extract_frame_png_returns_a_png(dataset: str) -> None:
+ from lelab.episode_media import extract_frame_png, locate_episode_video, resolve_dataset_dir
+
+ loc = locate_episode_video(resolve_dataset_dir(dataset), 0, camera="top")
+ png = extract_frame_png(loc, 3)
+ assert png[:4] == b"\x89PNG"
+
+
+def test_extract_frame_png_respects_max_width(dataset: str) -> None:
+ import io
+
+ from PIL import Image
+
+ from lelab.episode_media import extract_frame_png, locate_episode_video, resolve_dataset_dir
+
+ loc = locate_episode_video(resolve_dataset_dir(dataset), 0, camera="top")
+ img = Image.open(io.BytesIO(extract_frame_png(loc, 0, max_width=16)))
+ assert img.width == 16
+ assert img.height == 12 # 32x24 halved, aspect preserved
+
+
+def test_extract_frame_decodes_inside_the_right_episode(dataset: str) -> None:
+ """Frame 0 of episode 1 must not be frame 0 of the file.
+
+ Each frame carries a distinct grey level, so decoding the wrong window
+ yields a visibly different pixel — this is the check that the from_timestamp
+ offset is actually applied.
+ """
+ import io
+
+ from PIL import Image
+
+ from lelab.episode_media import extract_frame_png, locate_episode_video, resolve_dataset_dir
+
+ d = resolve_dataset_dir(dataset)
+ ep0 = Image.open(io.BytesIO(extract_frame_png(locate_episode_video(d, 0, camera="top"), 0)))
+ ep1 = Image.open(io.BytesIO(extract_frame_png(locate_episode_video(d, 1, camera="top"), 0)))
+ # Episode 1 starts 12 frames into the file; its first frame is a later ramp
+ # value than episode 0's. Lossy encoding blurs exact values, so compare with
+ # tolerance rather than equality.
+ assert abs(ep0.getpixel((16, 12))[0] - ep1.getpixel((16, 12))[0]) > 10
+
+
+def test_extract_thumbnails_one_pass(dataset: str) -> None:
+ from lelab.episode_media import extract_thumbnails, locate_episode_video, resolve_dataset_dir
+
+ loc = locate_episode_video(resolve_dataset_dir(dataset), 0, camera="top")
+ thumbs = extract_thumbnails(loc, [0, 4, 8], max_width=16)
+ assert len(thumbs) == 3
+ assert all(t[:4] == b"\x89PNG" for t in thumbs)
+
+
+def test_extract_thumbnails_empty_request(dataset: str) -> None:
+ from lelab.episode_media import extract_thumbnails, locate_episode_video, resolve_dataset_dir
+
+ loc = locate_episode_video(resolve_dataset_dir(dataset), 0, camera="top")
+ assert extract_thumbnails(loc, []) == []
+
+
+# ── action stream / motion ──────────────────────────────────────────────────
+
+
+def test_load_episode_actions_filters_to_the_episode(dataset: str) -> None:
+ from lelab.episode_media import load_episode_actions, resolve_dataset_dir
+
+ d = resolve_dataset_dir(dataset)
+ assert len(load_episode_actions(d, 0)) == 12
+ assert len(load_episode_actions(d, 1)) == 8
+ assert load_episode_actions(d, 1)[0] == [0.0, 0.0, 0.0]
+
+
+def test_motion_trace_shape_and_values(dataset: str) -> None:
+ from lelab.episode_media import episode_motion_trace, resolve_dataset_dir
+
+ trace = episode_motion_trace(resolve_dataset_dir(dataset), 0)
+ assert len(trace) == 12 # one value per frame
+ assert trace[0] == 0.0 # frame 0 has no predecessor
+ # Each joint steps by 1 per frame across 3 joints.
+ assert trace[1] == pytest.approx(3.0)
+
+
+def test_motion_trace_missing_episode(dataset: str) -> None:
+ from lelab.episode_media import EpisodeNotFoundError, episode_motion_trace, resolve_dataset_dir
+
+ with pytest.raises(EpisodeNotFoundError):
+ episode_motion_trace(resolve_dataset_dir(dataset), 99)