Add dataset record liveview
This commit is contained in:
parent
b5b2dbd7a0
commit
aa4242145c
1
.gitignore
vendored
1
.gitignore
vendored
@ -1,6 +1,7 @@
|
||||
# ============================
|
||||
# Python
|
||||
# ============================
|
||||
**/logs/**
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$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.
|
||||
|
||||
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.
|
||||
|
||||
@ -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"
|
||||
|
||||
@ -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:
|
||||
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
|
||||
|
||||
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