feat: 对齐ACT双相机帧

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)

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"