Add dataset record liveview

This commit is contained in:
ChenYuhan 2026-08-18 10:01:07 +08:00
parent b5b2dbd7a0
commit aa4242145c
6 changed files with 544 additions and 6 deletions

1
.gitignore vendored
View File

@ -1,6 +1,7 @@
# ============================
# Python
# ============================
**/logs/**
__pycache__/
*.py[cod]
*$py.class

View File

@ -181,6 +181,17 @@ uv run record --config_path config/gello/xarm7_gello_record_config.yaml -a
Controls: `Space` start the episode, `→` save it, `←` discard and re-record it, `Esc` stop recording. The arm resets to its initial point between episodes.
The xArm7 GELLO example also enables a lightweight browser camera preview at
`http://127.0.0.1:8765/` (or `http://<recorder-ip>:8765/` from another machine).
It reuses frames already captured by the recorder: JPEG encoding and HTTP
streaming run on background threads, the preview is capped at 8 FPS, and stale
preview frames are dropped rather than delaying dataset collection. Disable it
or tune it in the top-level `web_preview` section of the YAML configuration.
When `robot.enable_logs` is enabled for diagnostics, one
`logs/gello_record_sync_*.csv` file is written per episode. The
`preview_clients`, `record_period_ms`, `frame_overrun_ms`, `action_age_ms`, and
`preview_publish_ms` columns can be used to measure preview timing impact.
> During collection the **relative position between the robot arm and the camera must not change**, and the camera setup at inference time must match the one used during collection. If the arm or camera moves, previously collected data becomes invalid.
> If the dataset root already exists and `-r` is not given, the script asks whether to overwrite it, resume, or cancel.

View File

@ -51,6 +51,17 @@ teleop:
gripper_open_deg: 198.28125
gripper_close_deg: 155.75
# Lossy browser preview of the frames already captured for the dataset.
# JPEG encoding and network I/O run in the background; preview frames are
# dropped instead of delaying the 30 Hz recorder or 60 Hz GELLO control loop.
web_preview:
enabled: true
host: "0.0.0.0"
port: 8765
fps: 8
width: 480
jpeg_quality: 65
dataset:
# Dataset path relative to the directory where the command is started.
root: "datasets/xarm7_gello_joint_safe_datas"

View File

@ -7,7 +7,7 @@ import argparse
import logging
import shutil
import threading
from dataclasses import dataclass
from dataclasses import dataclass, field
from pathlib import Path
import lerobot_robot_ufactory # patch
from lerobot.scripts.lerobot_record import *
@ -16,13 +16,17 @@ from lerobot_robot_ufactory.teleoperators.uf_mock_teleop import UFMockTeleop
from lerobot_robot_ufactory.teleoperators.base_teleop import UFBaseTeleop
from lerobot_robot_ufactory.utils.realtime_teleop import RealtimeTeleopController
from lerobot_robot_ufactory.utils.utils import init_keyboard_listener
from lerobot_robot_ufactory.utils.web_preview import RecordingWebPreview, WebPreviewConfig
@dataclass
class UFRecordConfig(LeRobotRecordConfig):
"""RecordConfig variant that permits UFACTORY manual-mode recording."""
web_preview: WebPreviewConfig = field(default_factory=WebPreviewConfig)
def __post_init__(self):
self.web_preview.validate()
manual_mode = getattr(self.robot, "manual_mode", False)
if manual_mode:
if self.teleop is not None or self.policy is not None:
@ -276,19 +280,24 @@ def _disconnect_recording_resources(robot, teleop, listener):
class _RecordingCleanup:
def __init__(self, robot, teleop, listener, async_episode_saver):
def __init__(self, robot, teleop, listener, async_episode_saver, web_preview=None):
self.robot = robot
self.teleop = teleop
self.listener = listener
self.async_episode_saver = async_episode_saver
self.web_preview = web_preview
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
try:
if self.async_episode_saver is not None:
self.async_episode_saver.close()
try:
if self.async_episode_saver is not None:
self.async_episode_saver.close()
finally:
if self.web_preview is not None:
self.web_preview.stop()
finally:
_disconnect_recording_resources(self.robot, self.teleop, self.listener)
return False
@ -321,6 +330,7 @@ def record_loop(
manual_mode: bool = False,
manual_gripper_keys: dict[str, bool] | None = None,
manual_gripper_speed: float = 0.5,
web_preview: RecordingWebPreview | None = None,
):
if dataset is not None and dataset.fps != fps:
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")
@ -402,7 +412,16 @@ def record_loop(
"observation_end_s",
"state_to_observation_end_ms",
"camera_timings",
"preview_publish_ms",
"preview_clients",
"preview_source_generation",
"preview_encoded_frames",
"preview_last_encode_ms",
"preview_max_encode_ms",
"record_period_ms",
"frame_loop_ms",
"frame_budget_ms",
"frame_overrun_ms",
],
)
sync_log_writer.writeheader()
@ -410,8 +429,15 @@ def record_loop(
timestamp = 0
start_episode_t = time.perf_counter()
previous_loop_start_t = None
while timestamp < control_time_s:
start_loop_t = time.perf_counter()
record_period_ms = (
0.0
if previous_loop_start_t is None
else (start_loop_t - previous_loop_start_t) * 1000
)
previous_loop_start_t = start_loop_t
if events["exit_early"]:
events["exit_early"] = False
@ -435,6 +461,13 @@ def record_loop(
# Applies a pipeline to the raw robot observation, default is IdentityProcessor
obs_processed = robot_observation_processor(obs)
preview_publish_ms = 0.0
if web_preview is not None:
# This only replaces references in a latest-frame slot. All image
# processing and network I/O remain on preview background threads.
before_preview_publish_t = time.perf_counter()
web_preview.publish(obs_processed)
preview_publish_ms = (time.perf_counter() - before_preview_publish_t) * 1000
if policy is not None or dataset is not None:
observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR)
@ -542,6 +575,9 @@ def record_loop(
robot, "_last_realtime_observation_end_monotonic_s", observation_monotonic_s
)
camera_timings = getattr(robot, "_last_realtime_camera_timings", {})
preview_stats = web_preview.timing_stats() if web_preview is not None else {}
frame_loop_ms = (time.perf_counter() - start_loop_t) * 1000
frame_budget_ms = 1000 / fps
sync_log_writer.writerow(
{
"frame": sync_frame_index,
@ -551,7 +587,18 @@ def record_loop(
"observation_end_s": f"{observation_end_s:.9f}",
"state_to_observation_end_ms": f"{(observation_end_s - observation_monotonic_s) * 1000:.3f}",
"camera_timings": repr(camera_timings),
"frame_loop_ms": f"{(time.perf_counter() - start_loop_t) * 1000:.3f}",
"preview_publish_ms": f"{preview_publish_ms:.6f}",
"preview_clients": preview_stats.get("preview_clients", 0),
"preview_source_generation": preview_stats.get(
"preview_source_generation", 0
),
"preview_encoded_frames": preview_stats.get("preview_encoded_frames", 0),
"preview_last_encode_ms": f'{preview_stats.get("preview_last_encode_ms", 0.0):.3f}',
"preview_max_encode_ms": f'{preview_stats.get("preview_max_encode_ms", 0.0):.3f}',
"record_period_ms": f"{record_period_ms:.3f}",
"frame_loop_ms": f"{frame_loop_ms:.3f}",
"frame_budget_ms": f"{frame_budget_ms:.3f}",
"frame_overrun_ms": f"{max(0.0, frame_loop_ms - frame_budget_ms):.3f}",
}
)
sync_frame_index += 1
@ -759,12 +806,23 @@ def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset:
},
)
web_preview = None
try:
robot.connect()
if teleop is not None:
teleop.connect()
if cfg.web_preview.enabled:
web_preview = RecordingWebPreview(cfg.web_preview)
web_preview.start()
print(f"Camera web preview: {web_preview.url}")
if cfg.web_preview.host == "0.0.0.0":
print(
f"From another machine: http://<recorder-ip>:{cfg.web_preview.port}/"
)
except BaseException:
try:
if web_preview is not None:
web_preview.stop()
_disconnect_recording_resources(robot, teleop, None)
except BaseException:
logging.exception("Failed to clean up after recording device connection failure")
@ -829,7 +887,9 @@ def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset:
if async_episode_saver is not None:
print('Async episode saving is enabled.')
with _RecordingCleanup(robot, teleop, listener, async_episode_saver), VideoEncodingManager(dataset):
with _RecordingCleanup(
robot, teleop, listener, async_episode_saver, web_preview
), VideoEncodingManager(dataset):
recorded_episodes = 0
while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
time.sleep(0.01)
@ -875,6 +935,7 @@ def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset:
manual_mode=manual_mode,
manual_gripper_keys=manual_gripper_keys,
manual_gripper_speed=getattr(cfg.robot, "manual_gripper_speed", 0.5),
web_preview=web_preview,
)
else:
continue

View File

@ -0,0 +1,332 @@
"""Best-effort web preview for frames already captured by a recording loop.
The recorder only publishes references to its latest frames. Resizing, JPEG
encoding, and network writes happen on background threads, and preview frames
are deliberately dropped whenever those threads cannot keep up.
"""
from __future__ import annotations
import http.server
import json
import logging
import threading
import time
from dataclasses import dataclass
from typing import Any
from urllib.parse import parse_qs, urlsplit
import cv2
import numpy as np
logger = logging.getLogger(__name__)
@dataclass
class WebPreviewConfig:
"""Configuration for the optional recording-time camera preview."""
enabled: bool = False
host: str = "0.0.0.0"
port: int = 8765
fps: float = 8.0
width: int = 480
jpeg_quality: int = 65
def validate(self) -> None:
if not self.host:
raise ValueError("web_preview.host must not be empty")
if not 1 <= self.port <= 65535:
raise ValueError("web_preview.port must be between 1 and 65535")
if self.fps <= 0:
raise ValueError("web_preview.fps must be positive")
if self.width <= 0:
raise ValueError("web_preview.width must be positive")
if not 1 <= self.jpeg_quality <= 100:
raise ValueError("web_preview.jpeg_quality must be between 1 and 100")
class _PreviewFeed:
def __init__(self, name: str) -> None:
self.name = name
self.condition = threading.Condition()
self.jpeg: bytes | None = None
self.frame_id = 0
self.clients = 0
self.stopped = False
def add_client(self) -> None:
with self.condition:
self.clients += 1
def remove_client(self) -> None:
with self.condition:
self.clients = max(0, self.clients - 1)
def publish_jpeg(self, jpeg: bytes) -> None:
with self.condition:
self.jpeg = jpeg
self.frame_id += 1
self.condition.notify_all()
def wait_for_jpeg(self, previous_id: int) -> tuple[bytes | None, int, bool]:
with self.condition:
self.condition.wait_for(
lambda: self.frame_id > previous_id or self.stopped,
timeout=1.0,
)
return self.jpeg, self.frame_id, self.stopped
def stop(self) -> None:
with self.condition:
self.stopped = True
self.condition.notify_all()
class RecordingWebPreview:
"""Serve a lossy, asynchronous preview of recording observations."""
def __init__(self, config: WebPreviewConfig) -> None:
config.validate()
self.config = config
self._condition = threading.Condition()
self._latest_frames: dict[str, np.ndarray] = {}
self._source_generation = 0
self._encoded_frames = 0
self._last_encode_ms = 0.0
self._max_encode_ms = 0.0
self._stop_event = threading.Event()
self._feeds: dict[str, _PreviewFeed] = {}
self._encoder_thread: threading.Thread | None = None
self._server: http.server.ThreadingHTTPServer | None = None
self._server_thread: threading.Thread | None = None
@property
def url(self) -> str:
host = "127.0.0.1" if self.config.host == "0.0.0.0" else self.config.host
return f"http://{host}:{self.config.port}/"
def start(self) -> None:
if self._server is not None:
return
handler = type("RecordingPreviewHandler", (_PreviewHandler,), {"preview": self})
self._server = http.server.ThreadingHTTPServer(
(self.config.host, self.config.port), handler
)
self._server.daemon_threads = True
self._encoder_thread = threading.Thread(
target=self._encode_loop,
name="uf-recording-preview-encoder",
daemon=True,
)
self._server_thread = threading.Thread(
target=self._server.serve_forever,
name="uf-recording-preview-http",
daemon=True,
)
self._encoder_thread.start()
self._server_thread.start()
def publish(self, observation: dict[str, Any]) -> None:
"""Non-blockingly replace the latest previewable image references."""
frames = {
key: value
for key, value in observation.items()
if isinstance(value, np.ndarray) and value.ndim == 3 and value.shape[2] in (3, 4)
}
if not frames:
return
with self._condition:
self._latest_frames = frames
for name in frames:
self._feeds.setdefault(name, _PreviewFeed(name))
self._source_generation += 1
self._condition.notify()
def camera_names(self) -> list[str]:
with self._condition:
return sorted(self._feeds)
def feed(self, name: str | None) -> _PreviewFeed | None:
if name is None:
return None
with self._condition:
return self._feeds.get(name)
def timing_stats(self) -> dict[str, float | int]:
"""Return a cheap snapshot for recording synchronization diagnostics."""
with self._condition:
return {
"preview_clients": sum(feed.clients for feed in self._feeds.values()),
"preview_source_generation": self._source_generation,
"preview_encoded_frames": self._encoded_frames,
"preview_last_encode_ms": self._last_encode_ms,
"preview_max_encode_ms": self._max_encode_ms,
}
def _has_clients(self) -> bool:
return any(feed.clients > 0 for feed in self._feeds.values())
def _encode_loop(self) -> None:
last_generation = 0
next_encode_at = 0.0
period = 1.0 / self.config.fps
while not self._stop_event.is_set():
with self._condition:
self._condition.wait_for(
lambda: self._stop_event.is_set()
or (
self._source_generation > last_generation
and self._has_clients()
),
timeout=1.0,
)
if self._stop_event.is_set():
return
generation = self._source_generation
frames = dict(self._latest_frames)
delay = next_encode_at - time.monotonic()
if delay > 0 and self._stop_event.wait(delay):
return
next_encode_at = time.monotonic() + period
last_generation = generation
for name, frame in frames.items():
feed = self.feed(name)
if feed is None or feed.clients == 0:
continue
try:
encode_started = time.perf_counter()
jpeg = self._encode_frame(frame)
encode_ms = (time.perf_counter() - encode_started) * 1000
except Exception:
logger.exception("Failed to encode web preview frame for %s", name)
continue
feed.publish_jpeg(jpeg)
with self._condition:
self._encoded_frames += 1
self._last_encode_ms = encode_ms
self._max_encode_ms = max(self._max_encode_ms, encode_ms)
def _encode_frame(self, frame: np.ndarray) -> bytes:
height, width = frame.shape[:2]
if width != self.config.width:
target_height = max(1, round(height * self.config.width / width))
frame = cv2.resize(
frame,
(self.config.width, target_height),
interpolation=cv2.INTER_AREA,
)
if frame.shape[2] == 4:
frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
else:
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
ok, buffer = cv2.imencode(
".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, self.config.jpeg_quality]
)
if not ok:
raise RuntimeError("JPEG encoding failed")
return buffer.tobytes()
def stop(self) -> None:
if self._stop_event.is_set():
return
self._stop_event.set()
with self._condition:
self._condition.notify_all()
for feed in list(self._feeds.values()):
feed.stop()
if self._server is not None:
self._server.shutdown()
self._server.server_close()
if self._server_thread is not None:
self._server_thread.join(timeout=2.0)
if self._encoder_thread is not None:
self._encoder_thread.join(timeout=2.0)
_WEB_PAGE = """<!doctype html>
<html lang="zh-CN"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>LeRobot 录制预览</title>
<style>
:root{color-scheme:dark;font-family:system-ui,sans-serif}body{margin:0;background:#101214;color:#eee}
header{padding:14px 18px;background:#191c20;border-bottom:1px solid #333}h1{margin:0;font-size:20px}
#status{margin-top:7px;color:#aeb5bd;font-size:13px}.grid{display:grid;grid-template-columns:repeat(auto-fit,minmax(320px,1fr));gap:12px;padding:14px}
.tile{overflow:hidden;background:#191c20;border:1px solid #333;border-radius:7px}.tile h2{margin:0;padding:9px 11px;font-size:14px}.tile img{display:block;width:100%;background:#08090a}
</style></head><body><header><h1>LeRobot 录制相机预览</h1><div id="status">正在等待相机帧</div></header>
<main id="grid" class="grid"></main><script>
let signature='';async function refresh(){try{const r=await fetch('/api/status',{cache:'no-store'});const p=await r.json();
document.getElementById('status').textContent=p.cameras.length?`${p.cameras.length} 路相机 · 预览 ${p.fps} FPS · 采集优先预览允许丢帧`:'正在等待相机帧…';
const next=p.cameras.join('|');if(next!==signature){signature=next;const grid=document.getElementById('grid');grid.replaceChildren();
for(const name of p.cameras){const tile=document.createElement('section');tile.className='tile';const h=document.createElement('h2');h.textContent=name;
const img=document.createElement('img');img.alt=name;img.src='/stream?camera='+encodeURIComponent(name);tile.append(h,img);grid.appendChild(tile);}}}catch(e){document.getElementById('status').textContent='预览服务不可用:'+e}}
refresh();setInterval(refresh,2000);</script></body></html>""".encode()
class _PreviewHandler(http.server.BaseHTTPRequestHandler):
preview: RecordingWebPreview
def do_GET(self) -> None:
parsed = urlsplit(self.path)
if parsed.path == "/":
self._send(200, "text/html; charset=utf-8", _WEB_PAGE)
elif parsed.path == "/api/status":
body = json.dumps(
{"ok": True, "cameras": self.preview.camera_names(), "fps": self.preview.config.fps}
).encode()
self._send(200, "application/json; charset=utf-8", body)
elif parsed.path == "/stream":
self._stream(parse_qs(parsed.query).get("camera", [None])[0])
elif parsed.path == "/favicon.ico":
self._send(204, "text/plain", b"")
else:
self.send_error(404)
def _send(self, status: int, content_type: str, body: bytes) -> None:
try:
self.send_response(status)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
except (BrokenPipeError, ConnectionResetError):
pass
def _stream(self, name: str | None) -> None:
feed = self.preview.feed(name)
if feed is None:
self._send(404, "text/plain; charset=utf-8", b"Unknown camera\n")
return
feed.add_client()
with self.preview._condition:
self.preview._condition.notify()
try:
self.send_response(200)
self.send_header("Cache-Control", "no-store")
self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
self.end_headers()
frame_id = 0
while True:
jpeg, next_id, stopped = feed.wait_for_jpeg(frame_id)
if stopped:
return
if jpeg is None or next_id == frame_id:
continue
frame_id = next_id
self.wfile.write(
b"--frame\r\nContent-Type: image/jpeg\r\n"
+ f"Content-Length: {len(jpeg)}\r\n\r\n".encode()
+ jpeg
+ b"\r\n"
)
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass
finally:
feed.remove_client()
def log_message(self, format: str, *args: Any) -> None:
pass

122
tests/test_web_preview.py Normal file
View File

@ -0,0 +1,122 @@
import socket
import threading
import time
from pathlib import Path
from urllib.request import urlopen
import numpy as np
import pytest
import yaml
from lerobot_robot_ufactory.utils.web_preview import (
RecordingWebPreview,
WebPreviewConfig,
_WEB_PAGE,
)
def _preview(**overrides):
values = {
"enabled": True,
"host": "127.0.0.1",
"port": 8765,
"fps": 30,
"width": 32,
"jpeg_quality": 60,
}
values.update(overrides)
return RecordingWebPreview(WebPreviewConfig(**values))
def _start_encoder(preview):
preview._encoder_thread = threading.Thread(target=preview._encode_loop, daemon=True)
preview._encoder_thread.start()
def test_config_rejects_invalid_resource_limits():
with pytest.raises(ValueError, match="fps"):
WebPreviewConfig(fps=0).validate()
with pytest.raises(ValueError, match="jpeg_quality"):
WebPreviewConfig(jpeg_quality=101).validate()
def test_embedded_page_has_parseable_camera_signature_expression():
page = _WEB_PAGE.decode()
assert "const next=p.cameras.join('|')" in page
assert "join('\n')" not in page
def test_gello_record_config_enables_low_rate_web_preview():
config_path = Path("config/gello/xarm7_gello_record_config.yaml")
config = yaml.safe_load(config_path.read_text())
assert config["web_preview"]["enabled"] is True
assert config["web_preview"]["fps"] == 8
assert config["web_preview"]["width"] == 480
def test_publish_keeps_only_latest_frame_references():
preview = _preview()
first = np.zeros((24, 32, 3), dtype=np.uint8)
second = np.ones((24, 32, 3), dtype=np.uint8)
preview.publish({"J1.pos": 0.0, "camera": first})
preview.publish({"camera": second})
assert preview.camera_names() == ["camera"]
assert preview._latest_frames["camera"] is second
assert preview._source_generation == 2
def test_encoder_is_idle_without_a_browser_client(monkeypatch):
preview = _preview()
calls = []
monkeypatch.setattr(preview, "_encode_frame", lambda frame: calls.append(frame) or b"jpeg")
_start_encoder(preview)
try:
preview.publish({"camera": np.zeros((24, 32, 3), dtype=np.uint8)})
time.sleep(0.08)
assert calls == []
finally:
preview.stop()
def test_client_receives_background_encoded_latest_frame():
preview = _preview()
_start_encoder(preview)
try:
frame = np.zeros((24, 32, 3), dtype=np.uint8)
frame[:, :, 0] = 255
preview.publish({"camera": frame})
feed = preview.feed("camera")
feed.add_client()
with preview._condition:
preview._condition.notify()
jpeg, frame_id, stopped = feed.wait_for_jpeg(0)
feed.remove_client()
assert jpeg is not None and jpeg.startswith(b"\xff\xd8")
assert frame_id == 1
assert not stopped
finally:
preview.stop()
def test_status_endpoint_lists_published_cameras():
try:
with socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
port = sock.getsockname()[1]
except PermissionError:
pytest.skip("local sockets are disabled by the test sandbox")
preview = _preview(port=port)
preview.start()
try:
preview.publish({"camera2": np.zeros((24, 32, 3), dtype=np.uint8)})
with urlopen(preview.url + "api/status", timeout=2) as response:
body = response.read()
assert response.status == 200
assert b'"camera2"' in body
finally:
preview.stop()