diff --git a/.gitignore b/.gitignore index 3f7fe02..c802b7d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ # ============================ # Python # ============================ +**/logs/** __pycache__/ *.py[cod] *$py.class diff --git a/README.md b/README.md index ddcf5cc..cb683d2 100644 --- a/README.md +++ b/README.md @@ -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://: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. diff --git a/config/gello/xarm7_gello_record_config.yaml b/config/gello/xarm7_gello_record_config.yaml index e06531d..b693487 100644 --- a/config/gello/xarm7_gello_record_config.yaml +++ b/config/gello/xarm7_gello_record_config.yaml @@ -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" diff --git a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py index 72038fc..282683c 100644 --- a/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py +++ b/src/lerobot_robot_ufactory/scripts/uf_lerobot_record.py @@ -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://:{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 diff --git a/src/lerobot_robot_ufactory/utils/web_preview.py b/src/lerobot_robot_ufactory/utils/web_preview.py new file mode 100644 index 0000000..f48bbfe --- /dev/null +++ b/src/lerobot_robot_ufactory/utils/web_preview.py @@ -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 = """ + + +LeRobot 录制预览 +

LeRobot 录制相机预览

正在等待相机帧…
+
""".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 diff --git a/tests/test_web_preview.py b/tests/test_web_preview.py new file mode 100644 index 0000000..9bb105e --- /dev/null +++ b/tests/test_web_preview.py @@ -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()