From 1ec74f107ba8c7400b9a1075041ec498ed89d4cb Mon Sep 17 00:00:00 2001 From: YikaiFu-cart Date: Mon, 10 Aug 2026 18:06:45 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AF=B9=E9=BD=90ACT=E5=8F=8C=E7=9B=B8?= =?UTF-8?q?=E6=9C=BA=E5=B8=A7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../test/test_act_episode_recorder.py | 91 +++++++- .../xr_rm_teleop/act_episode_recorder.py | 199 ++++++++++++++++++ 2 files changed, 289 insertions(+), 1 deletion(-) diff --git a/xr_rm_teleop/test/test_act_episode_recorder.py b/xr_rm_teleop/test/test_act_episode_recorder.py index fda5988..cd9ba81 100644 --- a/xr_rm_teleop/test/test_act_episode_recorder.py +++ b/xr_rm_teleop/test/test_act_episode_recorder.py @@ -1,8 +1,16 @@ +import numpy as np +import pytest + from xr_rm_teleop.act_episode_recorder import ( NO_ACTION, ButtonTracker, + CameraBuffer, + CameraFrame, + QualityError, RecordingSession, RecordingState, + select_camera_pair, + select_frame, ) @@ -19,7 +27,7 @@ def _recording_session(*, origin_seq=10, max_samples=1800): return session -def test_recording_starts_on_first_sent_grip_action_and_samples_every_third_cycle(): +def test_recording_starts_on_first_sent_grip_action_and_downsamples(): session = RecordingSession(max_samples=1800) session.arm() @@ -184,3 +192,84 @@ def test_right_a_reports_recording_rejection_event(): ) assert events.reject_reason == "initial_pose_command_during_episode" + + +def _image(value, *, shape=(480, 640, 3), dtype=np.uint8): + return np.full(shape, value, dtype=dtype) + + +def test_select_frame_returns_latest_frame_not_after_control_time(): + frames = ( + CameraFrame(_image(1), 10, 100.0, 900_000_000), + CameraFrame(_image(2), 11, 133.3, 933_000_000), + CameraFrame(_image(3), 12, 166.6, 1_010_000_000), + ) + + selected, age_ms = select_frame(frames, 1_000_000_000, 100.0) + + assert selected.frame_number == 11 + assert age_ms == pytest.approx(67.0, abs=0.1) + + +def test_select_frame_rejects_missing_old_and_invalid_images(): + with pytest.raises(QualityError, match="camera_frame_missing"): + select_frame((), 1_000_000_000, 50.0) + with pytest.raises(QualityError, match="camera_frame_too_old"): + select_frame( + (CameraFrame(_image(1), 10, 100.0, 900_000_000),), + 1_000_000_000, + 50.0, + ) + with pytest.raises(QualityError, match="camera_frame_format"): + select_frame( + ( + CameraFrame( + _image(1, shape=(10, 10, 3)), + 10, + 100.0, + 990_000_000, + ), + ), + 1_000_000_000, + 50.0, + ) + + +def test_select_camera_pair_rejects_inter_camera_skew(): + high = (CameraFrame(_image(1), 10, 100.0, 990_000_000),) + wrist = (CameraFrame(_image(2), 20, 100.0, 930_000_000),) + + with pytest.raises(QualityError, match="camera_skew"): + select_camera_pair( + high, + wrist, + 1_000_000_000, + max_age_ms=100.0, + max_skew_ms=50.0, + ) + + +def test_camera_buffer_is_bounded_and_counts_dropped_frames(): + buffer = CameraBuffer(maxlen=4) + for frame_number in (10, 11, 13, 14, 15): + buffer.push( + CameraFrame( + _image(frame_number), + frame_number, + float(frame_number), + frame_number * 1_000_000, + ) + ) + + stats = buffer.stats() + + assert [frame.frame_number for frame in buffer.snapshot()] == [ + 11, + 13, + 14, + 15, + ] + assert stats.frame_count == 5 + assert stats.dropped_frames == 1 + assert stats.drop_ratio == pytest.approx(1.0 / 6.0) + assert stats.fps == pytest.approx(800.0) diff --git a/xr_rm_teleop/xr_rm_teleop/act_episode_recorder.py b/xr_rm_teleop/xr_rm_teleop/act_episode_recorder.py index 288060f..caffc2f 100644 --- a/xr_rm_teleop/xr_rm_teleop/act_episode_recorder.py +++ b/xr_rm_teleop/xr_rm_teleop/act_episode_recorder.py @@ -2,9 +2,208 @@ from __future__ import annotations +import threading +import time +from collections import deque from dataclasses import dataclass from enum import Enum +import numpy as np + + +class QualityError(RuntimeError): + pass + + +@dataclass(frozen=True) +class CameraFrame: + image: np.ndarray + frame_number: int + hardware_timestamp_ms: float + host_monotonic_ns: int + + +@dataclass(frozen=True) +class CameraStats: + frame_count: int + dropped_frames: int + first_host_monotonic_ns: int | None + last_host_monotonic_ns: int | None + + @property + def drop_ratio(self) -> float: + expected = self.frame_count + self.dropped_frames + return self.dropped_frames / expected if expected else 0.0 + + @property + def fps(self) -> float: + if self.frame_count < 2: + return 0.0 + assert self.first_host_monotonic_ns is not None + assert self.last_host_monotonic_ns is not None + elapsed_ns = self.last_host_monotonic_ns - self.first_host_monotonic_ns + if elapsed_ns <= 0: + return 0.0 + return (self.frame_count - 1) * 1e9 / elapsed_ns + + +class CameraBuffer: + def __init__(self, *, maxlen: int = 4) -> None: + if maxlen <= 0: + raise ValueError("maxlen must be positive") + self._frames: deque[CameraFrame] = deque(maxlen=maxlen) + self._lock = threading.Lock() + self._frame_count = 0 + self._dropped_frames = 0 + self._first_host_monotonic_ns: int | None = None + self._last_host_monotonic_ns: int | None = None + self._last_frame_number: int | None = None + + def push(self, frame: CameraFrame) -> None: + with self._lock: + if ( + self._last_frame_number is not None + and frame.frame_number > self._last_frame_number + 1 + ): + self._dropped_frames += ( + frame.frame_number - self._last_frame_number - 1 + ) + self._last_frame_number = frame.frame_number + self._frame_count += 1 + if self._first_host_monotonic_ns is None: + self._first_host_monotonic_ns = frame.host_monotonic_ns + self._last_host_monotonic_ns = frame.host_monotonic_ns + self._frames.append(frame) + + def snapshot(self) -> tuple[CameraFrame, ...]: + with self._lock: + return tuple(self._frames) + + def stats(self) -> CameraStats: + with self._lock: + return CameraStats( + frame_count=self._frame_count, + dropped_frames=self._dropped_frames, + first_host_monotonic_ns=self._first_host_monotonic_ns, + last_host_monotonic_ns=self._last_host_monotonic_ns, + ) + + +def select_frame( + frames: tuple[CameraFrame, ...], + control_monotonic_ns: int, + max_age_ms: float, +) -> tuple[CameraFrame, float]: + eligible = [ + frame + for frame in frames + if frame.host_monotonic_ns <= control_monotonic_ns + ] + if not eligible: + raise QualityError("camera_frame_missing") + frame = max(eligible, key=lambda item: item.host_monotonic_ns) + age_ms = (control_monotonic_ns - frame.host_monotonic_ns) * 1e-6 + if age_ms > max_age_ms: + raise QualityError("camera_frame_too_old") + if frame.image.shape != (480, 640, 3) or frame.image.dtype != np.uint8: + raise QualityError("camera_frame_format") + return frame, age_ms + + +def select_camera_pair( + high_frames: tuple[CameraFrame, ...], + wrist_frames: tuple[CameraFrame, ...], + control_monotonic_ns: int, + *, + max_age_ms: float, + max_skew_ms: float, +) -> tuple[CameraFrame, CameraFrame, float, float, float]: + high, high_age_ms = select_frame( + high_frames, control_monotonic_ns, max_age_ms + ) + wrist, wrist_age_ms = select_frame( + wrist_frames, control_monotonic_ns, max_age_ms + ) + skew_ms = abs(high.host_monotonic_ns - wrist.host_monotonic_ns) * 1e-6 + if skew_ms > max_skew_ms: + raise QualityError("camera_skew") + return high, wrist, high_age_ms, wrist_age_ms, skew_ms + + +class RealSenseCamera: + def __init__( + self, + serial: str, + expected_model: str, + *, + buffer_size: int = 4, + ) -> None: + self.serial = serial + self.expected_model = expected_model + self.buffer = CameraBuffer(maxlen=buffer_size) + self.last_error: Exception | None = None + self._stop = threading.Event() + self._thread: threading.Thread | None = None + self._pipeline = None + + def start(self) -> None: + import pyrealsense2 as rs + + if self._thread is not None: + raise RuntimeError("camera already started") + devices = { + device.get_info(rs.camera_info.serial_number): device + for device in rs.context().query_devices() + } + if self.serial not in devices: + raise RuntimeError(f"RealSense serial not found: {self.serial}") + model = devices[self.serial].get_info(rs.camera_info.name) + if self.expected_model not in model: + raise RuntimeError( + f"RealSense {self.serial} model mismatch: expected " + f"{self.expected_model}, got {model}" + ) + + pipeline = rs.pipeline() + config = rs.config() + config.enable_device(self.serial) + config.enable_stream(rs.stream.color, 640, 480, rs.format.rgb8, 30) + pipeline.start(config) + self._pipeline = pipeline + self._stop.clear() + self.last_error = None + self._thread = threading.Thread(target=self._capture, daemon=True) + self._thread.start() + + def stop(self) -> None: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=2.0) + if self._pipeline is not None: + self._pipeline.stop() + self._thread = None + self._pipeline = None + + def _capture(self) -> None: + try: + while not self._stop.is_set(): + frames = self._pipeline.wait_for_frames(timeout_ms=1000) + color = frames.get_color_frame() + received_ns = time.monotonic_ns() + if not color: + continue + self.buffer.push( + CameraFrame( + np.asanyarray(color.get_data()).copy(), + color.get_frame_number(), + color.get_timestamp(), + received_ns, + ) + ) + except Exception as exc: + if not self._stop.is_set(): + self.last_error = exc + class RecordingState(str, Enum): IDLE = "IDLE"