mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-07 10:01:56 +00:00
671 lines
28 KiB
Python
671 lines
28 KiB
Python
# Copyright 2024 The HuggingFace Inc. team. 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.
|
|
|
|
"""Foxglove visualization backend.
|
|
|
|
Live control-loop streaming (:func:`log_foxglove_data`) and seekable dataset playback
|
|
(:func:`serve_foxglove_dataset_playback`) over a Foxglove WebSocket server. Callers usually select a
|
|
backend at runtime through the dispatch in :mod:`lerobot.utils.visualization_utils` rather than
|
|
importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot[viz]'``).
|
|
"""
|
|
|
|
import logging
|
|
import numbers
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from lerobot.types import RobotAction, RobotObservation
|
|
|
|
from .constants import (
|
|
ACTION,
|
|
ACTION_PREFIX,
|
|
DONE,
|
|
OBS_IMAGES,
|
|
OBS_PREFIX,
|
|
OBS_STATE,
|
|
OBS_STR,
|
|
REWARD,
|
|
SUCCESS,
|
|
TRUNCATED,
|
|
)
|
|
from .import_utils import require_package
|
|
|
|
# Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}``
|
|
# pairs rather than one field per feature, so the same schema fits any robot regardless of which
|
|
# observation/action features it reports. The ``label`` field name is what Foxglove looks for to name
|
|
# each series automatically, so a single filtered path plots every feature, e.g.
|
|
# ``/observation/state.scalars[:]``.
|
|
_SCALARS_SCHEMA = {
|
|
"type": "object",
|
|
"title": "lerobot.Scalars",
|
|
"properties": {
|
|
"scalars": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"label": {"type": "string"},
|
|
"value": {"type": "number"},
|
|
},
|
|
},
|
|
}
|
|
},
|
|
}
|
|
|
|
|
|
def _is_scalar(x):
|
|
return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or (
|
|
isinstance(x, np.ndarray) and x.ndim == 0
|
|
)
|
|
|
|
|
|
def init_foxglove(host: str = "127.0.0.1", port: int | None = 8765) -> None:
|
|
"""
|
|
Starts a Foxglove WebSocket server for visualizing the control loop.
|
|
|
|
Connect to it from the Foxglove app at ``ws://<host>:<port>``. Calling this
|
|
more than once is a no-op while a server is already running.
|
|
|
|
Args:
|
|
host: Host interface to bind the WebSocket server to.
|
|
port: Port to bind the WebSocket server to (defaults to 8765).
|
|
"""
|
|
|
|
require_package("foxglove-sdk", extra="viz", import_name="foxglove")
|
|
import foxglove
|
|
|
|
# Live-stream state lives as attributes on ``log_foxglove_data``:
|
|
# ``.server`` is the shared WebSocket server and
|
|
# ``.channels`` caches one Foxglove channel per topic
|
|
if getattr(log_foxglove_data, "server", None) is not None:
|
|
return
|
|
log_foxglove_data.server = foxglove.start_server(host=host, port=port or 8765)
|
|
log_foxglove_data.channels = {}
|
|
|
|
|
|
def shutdown_foxglove() -> None:
|
|
"""Stops the Foxglove WebSocket server and clears cached channels."""
|
|
|
|
server = getattr(log_foxglove_data, "server", None)
|
|
if server is not None:
|
|
server.stop()
|
|
log_foxglove_data.server = None
|
|
log_foxglove_data.channels = {}
|
|
|
|
|
|
def _foxglove_safe_name(name: str) -> str:
|
|
"""Replace ``.`` with ``_`` so a feature name is a single Foxglove topic-path segment.
|
|
|
|
Foxglove treats ``.`` as a path separator, so an unsanitized name like ``observation.images.front``
|
|
would split into nested segments instead of naming one topic.
|
|
"""
|
|
|
|
return name.replace(".", "_")
|
|
|
|
|
|
def _foxglove_topic(key: str, *, is_image: bool = False) -> str:
|
|
"""Build the Foxglove topic for a feature ``key``.
|
|
|
|
Camera features map to a per-source image topic (``/observation/images/<name>``); scalar features
|
|
share one aggregate topic per source: ``/observation/state`` for observations, ``/action/state``
|
|
for actions.
|
|
"""
|
|
|
|
if is_image:
|
|
name = str(key)
|
|
for prefix in (f"{OBS_IMAGES}.", OBS_PREFIX):
|
|
if name.startswith(prefix):
|
|
name = name[len(prefix) :]
|
|
break
|
|
return f"/{OBS_STR}/images/{_foxglove_safe_name(name)}"
|
|
source = ACTION if (str(key).startswith(ACTION_PREFIX) or str(key) == ACTION) else OBS_STR
|
|
return f"/{source}/state"
|
|
|
|
|
|
def _log_foxglove_scalars(
|
|
topic: str, values: dict[str, float], *, channels: dict | None = None, log_time: int | None = None
|
|
) -> None:
|
|
"""Log scalars on a typed JSON channel using the static :data:`_SCALARS_SCHEMA`.
|
|
|
|
``values`` is an ordered mapping of feature name to value; it is emitted as a ``scalars`` array of
|
|
``{label, value}`` objects. Insertion order is preserved so series stay stable across messages.
|
|
|
|
``channels`` is the per-topic channel cache to reuse (defaults to the live-stream cache on
|
|
:func:`log_foxglove_data`; dataset playback passes its own local cache to stay self-contained).
|
|
``log_time`` is the message time in nanoseconds; when ``None`` the server's receive time is used.
|
|
"""
|
|
|
|
if not values:
|
|
return
|
|
|
|
import foxglove
|
|
|
|
if channels is None:
|
|
channels = log_foxglove_data.channels
|
|
channel = channels.get(topic)
|
|
if channel is None:
|
|
channel = channels[topic] = foxglove.Channel(topic, schema=_SCALARS_SCHEMA, message_encoding="json")
|
|
msg = {"scalars": [{"label": label, "value": value} for label, value in values.items()]}
|
|
if log_time is None:
|
|
channel.log(msg)
|
|
else:
|
|
channel.log(msg, log_time=log_time)
|
|
|
|
|
|
def _labeled_scalars(name: str, values, labels: list[str] | None = None) -> dict[str, float]:
|
|
"""Expand a 1D sequence into ``{label: value}`` entries with a consistent fallback."""
|
|
|
|
flat = [float(v) for v in values]
|
|
if labels is None or len(labels) != len(flat):
|
|
labels = [f"{name}_{i}" for i in range(len(flat))]
|
|
return dict(zip(labels, flat, strict=True))
|
|
|
|
|
|
def _log_foxglove_image(
|
|
topic: str,
|
|
frame_id: str,
|
|
arr: np.ndarray,
|
|
*,
|
|
compress_images: bool,
|
|
channels: dict | None = None,
|
|
log_time: int | None = None,
|
|
depth_range: tuple[float, float] | None = None,
|
|
) -> None:
|
|
"""Log an image on a cached per-topic channel.
|
|
|
|
Encoding is chosen from the frame's channel count and dtype:
|
|
|
|
- Single channel => depth map, logged uncompressed as ``16UC1`` (integer) or ``32FC1`` (float).
|
|
Plain ``uint8`` grayscale (with no ``depth_range``) is the exception and is sent as ``mono8``.
|
|
- Three channels => ``rgb8``.
|
|
- Four channels => ``rgba8``.
|
|
|
|
When ``compress_images`` is set, ``mono8`` and ``rgb8`` frames are JPEG-encoded instead; depth
|
|
and ``rgba8`` frames are always sent raw.
|
|
|
|
Args:
|
|
topic: Foxglove topic to log on.
|
|
frame_id: Frame id stamped on the message.
|
|
arr: Image as HWC or CHW (CHW is transposed to HWC), any dtype. Non-depth floating-point
|
|
images are assumed normalized to [0, 1] and scaled to uint8.
|
|
compress_images: JPEG-encode ``mono8`` and ``rgb8`` frames; ignored for depth and ``rgba8``.
|
|
channels: Per-topic channel cache to reuse (see :func:`_log_foxglove_scalars`).
|
|
log_time: Message time in nanoseconds, also written to the header timestamp; when ``None``
|
|
the server's receive time is used.
|
|
depth_range: ``(lo, hi)`` bounds that normalize depth to ``[0, 1]`` ``32FC1`` so Foxglove's
|
|
default panel scaling matches rerun's ``depth_range`` contrast (a ``RawImage`` carries
|
|
no value range itself).
|
|
"""
|
|
|
|
from foxglove.channels import CompressedImageChannel, RawImageChannel
|
|
from foxglove.messages import CompressedImage, RawImage, Timestamp
|
|
|
|
if channels is None:
|
|
channels = log_foxglove_data.channels
|
|
time_ns = time.time_ns() if log_time is None else log_time
|
|
timestamp = Timestamp(sec=time_ns // 1_000_000_000, nsec=time_ns % 1_000_000_000)
|
|
log_kwargs = {} if log_time is None else {"log_time": log_time}
|
|
|
|
# Convert CHW -> HWC when needed (mirrors log_rerun_data).
|
|
if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
|
|
arr = np.transpose(arr, (1, 2, 0))
|
|
height, width = arr.shape[0], arr.shape[1]
|
|
n_channels = 1 if arr.ndim == 2 else arr.shape[2]
|
|
|
|
# Single-channel => depth, unless it's plain uint8 grayscale (kept as mono8 below). Foxglove
|
|
# renders both 16UC1 (uint16) and 32FC1 (float32) depth with a colormap.
|
|
if n_channels == 1 and (depth_range is not None or arr.dtype != np.uint8):
|
|
depth = arr if arr.ndim == 2 else arr[..., 0]
|
|
if depth_range is not None:
|
|
lo, hi = depth_range
|
|
depth = depth.astype(np.float32)
|
|
depth = np.clip((depth - lo) / (hi - lo), 0.0, 1.0) if hi > lo else np.zeros_like(depth)
|
|
encoding, step = "32FC1", width * 4
|
|
elif np.issubdtype(depth.dtype, np.floating):
|
|
encoding, step = "32FC1", width * 4
|
|
else:
|
|
depth = np.clip(np.rint(depth), 0, 65535).astype("<u2")
|
|
encoding, step = "16UC1", width * 2
|
|
depth = np.ascontiguousarray(depth, dtype="<f4" if encoding == "32FC1" else "<u2")
|
|
channel = channels.get(topic)
|
|
if channel is None:
|
|
channel = channels[topic] = RawImageChannel(topic=topic)
|
|
channel.log(
|
|
RawImage(
|
|
timestamp=timestamp,
|
|
frame_id=frame_id,
|
|
width=width,
|
|
height=height,
|
|
encoding=encoding,
|
|
step=step,
|
|
data=depth.tobytes(),
|
|
),
|
|
**log_kwargs,
|
|
)
|
|
return
|
|
|
|
if np.issubdtype(arr.dtype, np.floating):
|
|
arr = (arr * 255.0).clip(0, 255)
|
|
arr = np.ascontiguousarray(arr, dtype=np.uint8)
|
|
|
|
if compress_images and n_channels in (1, 3):
|
|
buf_src = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) if n_channels == 3 else arr
|
|
_, buf = cv2.imencode(".jpg", buf_src)
|
|
channel = channels.get(topic)
|
|
if channel is None:
|
|
channel = channels[topic] = CompressedImageChannel(topic=topic)
|
|
channel.log(
|
|
CompressedImage(timestamp=timestamp, frame_id=frame_id, data=buf.tobytes(), format="jpeg"),
|
|
**log_kwargs,
|
|
)
|
|
return
|
|
|
|
encoding = {1: "mono8", 3: "rgb8", 4: "rgba8"}.get(n_channels)
|
|
if encoding is None:
|
|
logging.warning(
|
|
"Foxglove: skipping image on topic '%s' with unsupported shape %s (%d channels); "
|
|
"expected 1 (mono8), 3 (rgb8), or 4 (rgba8) channels.",
|
|
topic,
|
|
tuple(arr.shape),
|
|
n_channels,
|
|
)
|
|
return
|
|
channel = channels.get(topic)
|
|
if channel is None:
|
|
channel = channels[topic] = RawImageChannel(topic=topic)
|
|
channel.log(
|
|
RawImage(
|
|
timestamp=timestamp,
|
|
frame_id=frame_id,
|
|
width=width,
|
|
height=height,
|
|
encoding=encoding,
|
|
step=width * n_channels,
|
|
data=arr.tobytes(),
|
|
),
|
|
**log_kwargs,
|
|
)
|
|
|
|
|
|
def log_foxglove_data(
|
|
observation: RobotObservation | None = None,
|
|
action: RobotAction | None = None,
|
|
compress_images: bool = False,
|
|
) -> None:
|
|
"""
|
|
Logs observation and action data to a Foxglove WebSocket server for real-time visualization.
|
|
|
|
Mirrors ``log_rerun_data`` but emits Foxglove messages over the server started by
|
|
:func:`init_foxglove`. Data is mapped as follows:
|
|
- Scalars (and elements of 1D arrays) are accumulated per source and logged on the
|
|
``/observation/state`` and ``/action/state`` topics as typed JSON messages using the static
|
|
``lerobot.Scalars`` schema: a ``scalars`` array of ``{label, value}`` objects (see
|
|
:data:`_SCALARS_SCHEMA`). The ``label`` field lets Foxglove name each series automatically, so
|
|
``/observation/state.scalars[:].value`` plots every feature at once.
|
|
- 3D NumPy arrays that resemble images are transposed from CHW to HWC when needed and logged on a
|
|
per-source topic (e.g. ``/observation/images/front``) as a ``RawImage`` (or a JPEG
|
|
``CompressedImage`` when ``compress_images`` is True).
|
|
|
|
Args:
|
|
observation: An optional dictionary containing observation data to log.
|
|
action: An optional dictionary containing action data to log.
|
|
compress_images: Whether to JPEG-compress images before logging to save bandwidth in exchange
|
|
for CPU and quality.
|
|
"""
|
|
|
|
require_package("foxglove-sdk", extra="viz", import_name="foxglove")
|
|
|
|
if getattr(log_foxglove_data, "server", None) is None:
|
|
raise RuntimeError("init_foxglove() must be called before log_foxglove_data().")
|
|
|
|
now = time.time_ns()
|
|
|
|
if observation:
|
|
obs_scalars: dict[str, float] = {}
|
|
for k, v in observation.items():
|
|
if v is None:
|
|
continue
|
|
key = k[len(OBS_PREFIX) :] if str(k).startswith(OBS_PREFIX) else str(k)
|
|
if _is_scalar(v):
|
|
obs_scalars[key] = float(v)
|
|
elif isinstance(v, np.ndarray):
|
|
if v.ndim == 1:
|
|
obs_scalars.update(_labeled_scalars(key, v))
|
|
else:
|
|
_log_foxglove_image(
|
|
_foxglove_topic(k, is_image=True),
|
|
key,
|
|
v,
|
|
compress_images=compress_images,
|
|
log_time=now,
|
|
)
|
|
_log_foxglove_scalars(_foxglove_topic(OBS_STATE), obs_scalars, log_time=now)
|
|
|
|
if action:
|
|
action_scalars: dict[str, float] = {}
|
|
for k, v in action.items():
|
|
if v is None:
|
|
continue
|
|
key = k[len(ACTION_PREFIX) :] if str(k).startswith(ACTION_PREFIX) else str(k)
|
|
if _is_scalar(v):
|
|
action_scalars[key] = float(v)
|
|
elif isinstance(v, np.ndarray):
|
|
action_scalars.update(_labeled_scalars(key, v.flatten()))
|
|
_log_foxglove_scalars(_foxglove_topic(ACTION), action_scalars, log_time=now)
|
|
|
|
|
|
# ── Dataset playback over a Foxglove WebSocket server ─────────────────────
|
|
# A LeRobotDataset is random-access on disk, so rather than fire-and-forget a forward stream we
|
|
# advertise a seekable timeline and serve frames on demand for whatever time the user scrubs/plays
|
|
# to in the Foxglove app. This relies on the SDK's PlaybackControl capability.
|
|
|
|
|
|
def _feature_dim_names(feature: dict | None) -> list[str] | None:
|
|
"""Best-effort per-dimension series labels for a 1D feature, or ``None`` to fall back to indices.
|
|
|
|
LeRobot records a feature's ``names`` inconsistently: a flat list (``["x", "y"]``), a category
|
|
mapping (``{"motors": ["motor_0", "motor_1"]}``), or a name->index mapping
|
|
(``{"delta_x": 0, "delta_y": 1}``). Each is handled, but labels are only returned when their count
|
|
matches the feature's 1D shape, so a malformed/mismatched ``names`` can't silently mislabel series.
|
|
"""
|
|
|
|
if not feature:
|
|
return None
|
|
shape = feature.get("shape")
|
|
dim = shape[0] if shape and len(shape) == 1 else None
|
|
names = feature.get("names")
|
|
labels: list[str] | None = None
|
|
if isinstance(names, dict):
|
|
values = list(names.values())
|
|
if values and all(isinstance(v, (list, tuple)) for v in values):
|
|
labels = [str(n) for group in values for n in group]
|
|
elif values and all(isinstance(v, int) and not isinstance(v, bool) for v in values):
|
|
labels = [name for name, _ in sorted(names.items(), key=lambda kv: kv[1])]
|
|
elif isinstance(names, (list, tuple)):
|
|
labels = [str(n) for n in names]
|
|
if labels is not None and dim is not None and len(labels) == dim:
|
|
return labels
|
|
return None
|
|
|
|
|
|
def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -> dict[str, float]:
|
|
"""Flatten a frame's vector/scalar feature ``key`` into ``{label: value}`` entries.
|
|
|
|
``labels`` provides one name per dimension (from the dataset's feature metadata); when absent or
|
|
the wrong length, dimensions fall back to ``{name}_{i}`` (the short feature name), matching the
|
|
live stream so series names agree. A scalar feature becomes a single entry. Missing or ``None``
|
|
features yield an empty mapping.
|
|
"""
|
|
|
|
v = sample.get(key)
|
|
if v is None:
|
|
return {}
|
|
arr = v.numpy() if hasattr(v, "numpy") else np.asarray(v)
|
|
if key.startswith(OBS_PREFIX):
|
|
name = key[len(OBS_PREFIX) :]
|
|
elif key.startswith(ACTION_PREFIX):
|
|
name = key[len(ACTION_PREFIX) :]
|
|
else:
|
|
name = key
|
|
if arr.ndim == 0:
|
|
return {name: float(arr)}
|
|
return _labeled_scalars(name, arr.flatten(), labels)
|
|
|
|
|
|
def serve_foxglove_dataset_playback(
|
|
dataset,
|
|
episode_index: int,
|
|
*,
|
|
host: str = "127.0.0.1",
|
|
port: int = 8765,
|
|
compress_images: bool = False,
|
|
autoplay: bool = True,
|
|
) -> None:
|
|
"""Serve a single dataset episode to Foxglove as a seekable, scrubbable timeline.
|
|
|
|
Starts a Foxglove WebSocket server advertising the ``PlaybackControl`` capability over the
|
|
episode's time range. The Foxglove app drives play/pause/seek/speed; a background thread and a
|
|
``ServerListener`` read frames from the on-disk ``dataset`` on demand and log them stamped at
|
|
their dataset timestamps, so the user can scrub anywhere in the episode. Blocks until interrupted.
|
|
|
|
Args:
|
|
dataset: A ``LeRobotDataset`` loaded for the single episode to visualize.
|
|
episode_index: Index of the episode being visualized (used only for the session name).
|
|
host: Host interface to bind the WebSocket server to.
|
|
port: Port to bind the WebSocket server to.
|
|
compress_images: Whether to JPEG-compress camera frames before logging.
|
|
autoplay: If True, start playing automatically as soon as a client connects, instead of
|
|
waiting for the user to press play in the Foxglove app.
|
|
"""
|
|
|
|
require_package("foxglove-sdk", extra="viz", import_name="foxglove")
|
|
import bisect
|
|
import threading
|
|
|
|
import foxglove
|
|
from foxglove.websocket import (
|
|
Capability,
|
|
PlaybackCommand,
|
|
PlaybackControlRequest,
|
|
PlaybackState,
|
|
PlaybackStatus,
|
|
ServerListener,
|
|
)
|
|
|
|
# Per-frame timestamps in nanoseconds (read straight from the table, no video decode).
|
|
times_ns = [int(round(float(t) * 1e9)) for t in dataset.hf_dataset["timestamp"]]
|
|
n_frames = len(times_ns)
|
|
if n_frames == 0:
|
|
raise ValueError("Cannot visualize an empty episode.")
|
|
first_ns, last_ns = times_ns[0], times_ns[-1]
|
|
camera_keys = list(dataset.meta.camera_keys)
|
|
# Dataset-wide q01/q99 depth bounds (fallback min/max) used to normalize depth to [0, 1].
|
|
depth_ranges: dict[str, tuple[float, float]] = {}
|
|
for key in dataset.meta.depth_keys:
|
|
stats = (dataset.meta.stats or {}).get(key)
|
|
if not stats:
|
|
continue
|
|
lo = stats["q01"] if "q01" in stats else stats["min"]
|
|
hi = stats["q99"] if "q99" in stats else stats["max"]
|
|
depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
|
|
# Per-dimension series labels from the dataset metadata (e.g. joint names), computed once.
|
|
scalar_labels = {
|
|
OBS_STATE: _feature_dim_names(dataset.meta.features.get(OBS_STATE)),
|
|
ACTION: _feature_dim_names(dataset.meta.features.get(ACTION)),
|
|
}
|
|
# Local channel cache so the playback server is self-contained and doesn't touch the live-stream cache.
|
|
channels: dict = {}
|
|
|
|
def emit_frame(i: int) -> None:
|
|
"""Log every channel for frame ``i`` stamped at its dataset timestamp."""
|
|
sample = dataset[i]
|
|
log_time = times_ns[i]
|
|
for key in camera_keys:
|
|
arr = sample.get(key)
|
|
if arr is None:
|
|
continue
|
|
arr = arr.numpy() if hasattr(arr, "numpy") else np.asarray(arr)
|
|
_log_foxglove_image(
|
|
_foxglove_topic(key, is_image=True),
|
|
key,
|
|
arr,
|
|
compress_images=compress_images,
|
|
channels=channels,
|
|
log_time=log_time,
|
|
depth_range=depth_ranges.get(key),
|
|
)
|
|
_log_foxglove_scalars(
|
|
_foxglove_topic(OBS_STATE),
|
|
_frame_to_scalars(sample, OBS_STATE, scalar_labels[OBS_STATE]),
|
|
channels=channels,
|
|
log_time=log_time,
|
|
)
|
|
_log_foxglove_scalars(
|
|
_foxglove_topic(ACTION),
|
|
_frame_to_scalars(sample, ACTION, scalar_labels[ACTION]),
|
|
channels=channels,
|
|
log_time=log_time,
|
|
)
|
|
episode_scalars = {}
|
|
for feat, label in (
|
|
(DONE, "done"),
|
|
(TRUNCATED, "truncated"),
|
|
(REWARD, "reward"),
|
|
(SUCCESS, "success"),
|
|
):
|
|
v = sample.get(feat)
|
|
if v is not None:
|
|
episode_scalars[label] = float(v)
|
|
_log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time)
|
|
|
|
lock = threading.Lock()
|
|
stop_event = threading.Event()
|
|
# Shared playback state, guarded by ``lock``. ``seek_idx`` is a one-shot request set by the
|
|
# listener and serviced by the playback loop, which is the *only* thread that emits frames (so
|
|
# concurrent random access into the on-disk dataset / video decoder never overlaps).
|
|
state = {
|
|
"status": PlaybackStatus.Paused,
|
|
"cursor": first_ns,
|
|
"speed": 1.0,
|
|
"last_idx": -1,
|
|
"seek_idx": None,
|
|
}
|
|
|
|
def index_at(t_ns: int) -> int:
|
|
return max(0, min(n_frames - 1, bisect.bisect_right(times_ns, t_ns) - 1))
|
|
|
|
# One-shot latch so autoplay fires only on the first client subscription.
|
|
autoplay_started = threading.Event()
|
|
|
|
class _PlaybackListener(ServerListener):
|
|
def on_subscribe(self, client, channel):
|
|
# Start playing automatically once a client actually connects (subscribes). Using the
|
|
# subscribe hook, rather than starting in Playing up front, means the timeline doesn't
|
|
# advance before anyone is watching. Fires once; the user can still pause/seek after.
|
|
if not autoplay:
|
|
return
|
|
with lock:
|
|
if autoplay_started.is_set() or state["status"] != PlaybackStatus.Paused:
|
|
return
|
|
autoplay_started.set()
|
|
state["status"] = PlaybackStatus.Playing
|
|
cursor, speed = state["cursor"], state["speed"]
|
|
server.broadcast_playback_state(PlaybackState(PlaybackStatus.Playing, cursor, speed, False, ""))
|
|
|
|
def on_playback_control_request(self, req: PlaybackControlRequest):
|
|
# Only mutate state here; the playback loop performs all frame emission.
|
|
with lock:
|
|
did_seek = False
|
|
if req.seek_time is not None:
|
|
cursor = max(first_ns, min(last_ns, req.seek_time))
|
|
state["cursor"] = cursor
|
|
state["last_idx"] = state["seek_idx"] = index_at(cursor)
|
|
did_seek = True
|
|
if req.playback_speed and req.playback_speed > 0:
|
|
state["speed"] = req.playback_speed
|
|
if req.playback_command == PlaybackCommand.Play:
|
|
# Restarting from the end replays from the beginning.
|
|
if state["cursor"] >= last_ns:
|
|
state["cursor"] = first_ns
|
|
state["last_idx"] = state["seek_idx"] = 0
|
|
did_seek = True
|
|
state["status"] = PlaybackStatus.Playing
|
|
elif req.playback_command == PlaybackCommand.Pause:
|
|
state["status"] = PlaybackStatus.Paused
|
|
status, cursor, speed = state["status"], state["cursor"], state["speed"]
|
|
request_id = req.request_id or ""
|
|
return PlaybackState(status, cursor, speed, did_seek, request_id)
|
|
|
|
server = foxglove.start_server(
|
|
name=f"{dataset.repo_id}/episode_{episode_index}",
|
|
host=host,
|
|
port=port,
|
|
capabilities=[Capability.PlaybackControl, Capability.Time],
|
|
server_listener=_PlaybackListener(),
|
|
playback_time_range=(first_ns, last_ns),
|
|
)
|
|
|
|
def playback_loop() -> None:
|
|
# Cap how far the cursor may advance in a single tick. A slow frame decode (or any stall)
|
|
# would otherwise make ``dt`` huge and produce one enormous catch-up batch; clamping it makes
|
|
# playback trail wall-clock under a slow decoder while each tick emits a bounded frame range.
|
|
max_tick_dt_s = 0.25
|
|
prev = time.monotonic()
|
|
while not stop_event.is_set():
|
|
time.sleep(1.0 / 60.0)
|
|
ended = False
|
|
speed = 1.0
|
|
with lock:
|
|
now = time.monotonic()
|
|
dt = min(now - prev, max_tick_dt_s)
|
|
prev = now
|
|
# A queued seek is always serviced, even while paused, so scrubbing updates the view.
|
|
work = []
|
|
seek_idx = state["seek_idx"]
|
|
if seek_idx is not None:
|
|
state["seek_idx"] = None
|
|
work.append(seek_idx)
|
|
if state["status"] == PlaybackStatus.Playing:
|
|
cursor = state["cursor"] + int(dt * 1e9 * state["speed"])
|
|
start_idx = state["last_idx"] + 1
|
|
if cursor >= last_ns:
|
|
cursor, target, ended = last_ns, n_frames - 1, True
|
|
else:
|
|
target = index_at(cursor)
|
|
state["cursor"] = cursor
|
|
work.extend(range(start_idx, target + 1))
|
|
# cursor only grows while playing (seeks reset last_idx in the listener), so
|
|
# target >= last_idx here; a plain assignment is correct and clearer than max().
|
|
state["last_idx"] = target
|
|
if ended:
|
|
state["status"] = PlaybackStatus.Ended
|
|
if not work:
|
|
continue
|
|
cursor, speed = state["cursor"], state["speed"]
|
|
# Emit outside the lock; this is the only thread that calls emit_frame. Re-check
|
|
# stop_event between frames so shutdown stays responsive even mid-batch.
|
|
for i in work:
|
|
if stop_event.is_set():
|
|
break
|
|
emit_frame(i)
|
|
server.broadcast_time(cursor)
|
|
if ended:
|
|
server.broadcast_playback_state(PlaybackState(PlaybackStatus.Ended, cursor, speed, False, ""))
|
|
|
|
# Emit the first frame so channels are advertised (done before the loop starts, so emission stays
|
|
# single-threaded). Late-connecting clients re-receive frames once they seek/play.
|
|
emit_frame(0)
|
|
with lock:
|
|
state["last_idx"] = 0
|
|
server.broadcast_time(first_ns)
|
|
server.broadcast_playback_state(PlaybackState(PlaybackStatus.Paused, first_ns, 1.0, True, ""))
|
|
|
|
thread = threading.Thread(target=playback_loop, name="foxglove-playback", daemon=True)
|
|
thread.start()
|
|
|
|
print(f"Foxglove server running. Connect the Foxglove app to ws://{host}:{port}")
|
|
print("Use the playback controls in Foxglove to play/pause and scrub the episode. Ctrl-C to exit.")
|
|
try:
|
|
while not stop_event.is_set():
|
|
time.sleep(0.5)
|
|
except KeyboardInterrupt:
|
|
print("Ctrl-C received. Exiting.")
|
|
finally:
|
|
stop_event.set()
|
|
thread.join(timeout=2.0)
|
|
server.stop()
|
|
channels.clear()
|