Add dataset record liveview
This commit is contained in:
parent
b5b2dbd7a0
commit
aa4242145c
1
.gitignore
vendored
1
.gitignore
vendored
@ -1,6 +1,7 @@
|
|||||||
# ============================
|
# ============================
|
||||||
# Python
|
# Python
|
||||||
# ============================
|
# ============================
|
||||||
|
**/logs/**
|
||||||
__pycache__/
|
__pycache__/
|
||||||
*.py[cod]
|
*.py[cod]
|
||||||
*$py.class
|
*$py.class
|
||||||
|
|||||||
11
README.md
11
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.
|
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.
|
> 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.
|
> If the dataset root already exists and `-r` is not given, the script asks whether to overwrite it, resume, or cancel.
|
||||||
|
|||||||
@ -51,6 +51,17 @@ teleop:
|
|||||||
gripper_open_deg: 198.28125
|
gripper_open_deg: 198.28125
|
||||||
gripper_close_deg: 155.75
|
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:
|
||||||
# Dataset path relative to the directory where the command is started.
|
# Dataset path relative to the directory where the command is started.
|
||||||
root: "datasets/xarm7_gello_joint_safe_datas"
|
root: "datasets/xarm7_gello_joint_safe_datas"
|
||||||
|
|||||||
@ -7,7 +7,7 @@ import argparse
|
|||||||
import logging
|
import logging
|
||||||
import shutil
|
import shutil
|
||||||
import threading
|
import threading
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass, field
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import lerobot_robot_ufactory # patch
|
import lerobot_robot_ufactory # patch
|
||||||
from lerobot.scripts.lerobot_record import *
|
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.teleoperators.base_teleop import UFBaseTeleop
|
||||||
from lerobot_robot_ufactory.utils.realtime_teleop import RealtimeTeleopController
|
from lerobot_robot_ufactory.utils.realtime_teleop import RealtimeTeleopController
|
||||||
from lerobot_robot_ufactory.utils.utils import init_keyboard_listener
|
from lerobot_robot_ufactory.utils.utils import init_keyboard_listener
|
||||||
|
from lerobot_robot_ufactory.utils.web_preview import RecordingWebPreview, WebPreviewConfig
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class UFRecordConfig(LeRobotRecordConfig):
|
class UFRecordConfig(LeRobotRecordConfig):
|
||||||
"""RecordConfig variant that permits UFACTORY manual-mode recording."""
|
"""RecordConfig variant that permits UFACTORY manual-mode recording."""
|
||||||
|
|
||||||
|
web_preview: WebPreviewConfig = field(default_factory=WebPreviewConfig)
|
||||||
|
|
||||||
def __post_init__(self):
|
def __post_init__(self):
|
||||||
|
self.web_preview.validate()
|
||||||
manual_mode = getattr(self.robot, "manual_mode", False)
|
manual_mode = getattr(self.robot, "manual_mode", False)
|
||||||
if manual_mode:
|
if manual_mode:
|
||||||
if self.teleop is not None or self.policy is not None:
|
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:
|
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.robot = robot
|
||||||
self.teleop = teleop
|
self.teleop = teleop
|
||||||
self.listener = listener
|
self.listener = listener
|
||||||
self.async_episode_saver = async_episode_saver
|
self.async_episode_saver = async_episode_saver
|
||||||
|
self.web_preview = web_preview
|
||||||
|
|
||||||
def __enter__(self):
|
def __enter__(self):
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def __exit__(self, exc_type, exc_value, traceback):
|
def __exit__(self, exc_type, exc_value, traceback):
|
||||||
try:
|
try:
|
||||||
if self.async_episode_saver is not None:
|
try:
|
||||||
self.async_episode_saver.close()
|
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:
|
finally:
|
||||||
_disconnect_recording_resources(self.robot, self.teleop, self.listener)
|
_disconnect_recording_resources(self.robot, self.teleop, self.listener)
|
||||||
return False
|
return False
|
||||||
@ -321,6 +330,7 @@ def record_loop(
|
|||||||
manual_mode: bool = False,
|
manual_mode: bool = False,
|
||||||
manual_gripper_keys: dict[str, bool] | None = None,
|
manual_gripper_keys: dict[str, bool] | None = None,
|
||||||
manual_gripper_speed: float = 0.5,
|
manual_gripper_speed: float = 0.5,
|
||||||
|
web_preview: RecordingWebPreview | None = None,
|
||||||
):
|
):
|
||||||
if dataset is not None and dataset.fps != fps:
|
if dataset is not None and dataset.fps != fps:
|
||||||
raise ValueError(f"The dataset fps should be equal to requested fps ({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",
|
"observation_end_s",
|
||||||
"state_to_observation_end_ms",
|
"state_to_observation_end_ms",
|
||||||
"camera_timings",
|
"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_loop_ms",
|
||||||
|
"frame_budget_ms",
|
||||||
|
"frame_overrun_ms",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
sync_log_writer.writeheader()
|
sync_log_writer.writeheader()
|
||||||
@ -410,8 +429,15 @@ def record_loop(
|
|||||||
|
|
||||||
timestamp = 0
|
timestamp = 0
|
||||||
start_episode_t = time.perf_counter()
|
start_episode_t = time.perf_counter()
|
||||||
|
previous_loop_start_t = None
|
||||||
while timestamp < control_time_s:
|
while timestamp < control_time_s:
|
||||||
start_loop_t = time.perf_counter()
|
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"]:
|
if events["exit_early"]:
|
||||||
events["exit_early"] = False
|
events["exit_early"] = False
|
||||||
@ -435,6 +461,13 @@ def record_loop(
|
|||||||
|
|
||||||
# Applies a pipeline to the raw robot observation, default is IdentityProcessor
|
# Applies a pipeline to the raw robot observation, default is IdentityProcessor
|
||||||
obs_processed = robot_observation_processor(obs)
|
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:
|
if policy is not None or dataset is not None:
|
||||||
observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR)
|
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
|
robot, "_last_realtime_observation_end_monotonic_s", observation_monotonic_s
|
||||||
)
|
)
|
||||||
camera_timings = getattr(robot, "_last_realtime_camera_timings", {})
|
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(
|
sync_log_writer.writerow(
|
||||||
{
|
{
|
||||||
"frame": sync_frame_index,
|
"frame": sync_frame_index,
|
||||||
@ -551,7 +587,18 @@ def record_loop(
|
|||||||
"observation_end_s": f"{observation_end_s:.9f}",
|
"observation_end_s": f"{observation_end_s:.9f}",
|
||||||
"state_to_observation_end_ms": f"{(observation_end_s - observation_monotonic_s) * 1000:.3f}",
|
"state_to_observation_end_ms": f"{(observation_end_s - observation_monotonic_s) * 1000:.3f}",
|
||||||
"camera_timings": repr(camera_timings),
|
"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
|
sync_frame_index += 1
|
||||||
@ -759,12 +806,23 @@ def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset:
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
web_preview = None
|
||||||
try:
|
try:
|
||||||
robot.connect()
|
robot.connect()
|
||||||
if teleop is not None:
|
if teleop is not None:
|
||||||
teleop.connect()
|
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:
|
except BaseException:
|
||||||
try:
|
try:
|
||||||
|
if web_preview is not None:
|
||||||
|
web_preview.stop()
|
||||||
_disconnect_recording_resources(robot, teleop, None)
|
_disconnect_recording_resources(robot, teleop, None)
|
||||||
except BaseException:
|
except BaseException:
|
||||||
logging.exception("Failed to clean up after recording device connection failure")
|
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:
|
if async_episode_saver is not None:
|
||||||
print('Async episode saving is enabled.')
|
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
|
recorded_episodes = 0
|
||||||
while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
|
while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
|
||||||
time.sleep(0.01)
|
time.sleep(0.01)
|
||||||
@ -875,6 +935,7 @@ def record(cfg: UFRecordConfig, async_save: bool = False) -> LeRobotDataset:
|
|||||||
manual_mode=manual_mode,
|
manual_mode=manual_mode,
|
||||||
manual_gripper_keys=manual_gripper_keys,
|
manual_gripper_keys=manual_gripper_keys,
|
||||||
manual_gripper_speed=getattr(cfg.robot, "manual_gripper_speed", 0.5),
|
manual_gripper_speed=getattr(cfg.robot, "manual_gripper_speed", 0.5),
|
||||||
|
web_preview=web_preview,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
continue
|
continue
|
||||||
|
|||||||
332
src/lerobot_robot_ufactory/utils/web_preview.py
Normal file
332
src/lerobot_robot_ufactory/utils/web_preview.py
Normal 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
122
tests/test_web_preview.py
Normal 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()
|
||||||
Loading…
Reference in New Issue
Block a user