"""LiveRerunSession — canonical scaffold for live rerun overlay sessions. Bakes in the proven defaults from calibrate.py — these were debugged live on the YAM rig on 2026-06-02. Use this class for any recording, inference, or calibration script that needs: - A live ZED scene camera (HD1080 by default). - Optional teleop on right or left arm, with yellow-leader-button exit. - A rerun web viewer (gRPC + HTTP, JPEG-encoded logs). - A per-frame callback that gets the freshest BGR scene image and a monotonically-increasing frame index. Defaults that matter (each was a debugging milestone — don't change without re-validating): | Setting | Default | Why | |---|---|---| | ZED resolution | HD1080 | matches training data | | ZED grab strategy | single | ZED keeps only the latest frame; drain loops *block* per grab and ruin latency | | Rerun transport | gRPC + serve_web_viewer | matches raiden's working pattern | | Image encoding | JPEG via to_jpeg | raw HD1080 RGB at 15 fps swamps the gRPC channel (~280 MB/s); JPEG is ~50× lighter | | Rerun timeline | rr.set_time("step", sequence=frame_idx) | without this the viewer freezes on frame 0 | | Teleop exit signal | yellow top leader button | rc.check_button_press() polled at 20 Hz from the main thread | Render rate + scale are *user-controlled* (per-script), but the canonical pattern is **render every other frame, cache between**; ``calibrate.py`` shows the idiom. Example:: from lib.rerun_session import LiveRerunSession with LiveRerunSession("my_session", teleop_arm="right") as s: T_lock = s.lock_exo_pose(min_detections=10) W, H = s.WH cached_render = cached_mask = None def step(idx, bgr): nonlocal cached_render, cached_mask if idx % 2 == 0 or cached_render is None: joints7 = s.follower_joints7() rgb, mask = render_arm_with_exo(joints7, T_lock, s.K, W, H) cached_render = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) cached_mask = mask s.log("scene/rgb", bgr) s.log("scene/composite", compose_render_over_live(bgr, cached_render, cached_mask, 0.5)) s.loop(step) """ from __future__ import annotations import json import os import sys import threading import time from pathlib import Path from typing import Callable, Optional os.environ.setdefault("MUJOCO_GL", "egl") import cv2 import numpy as np import rerun as rr from .calibration import pool_exo_aruco from .rerun_helpers import init_web, to_jpeg from .teleop import start_teleop, stop_teleop CAM_CFG_PATH = Path.home() / ".config" / "raiden" / "camera.json" class LiveRerunSession: """Live ZED + rerun + (optional) teleop scaffold. Use as a context manager. Provides ``grab()``, ``follower_joints7()``, ``lock_exo_pose()``, ``log()``, ``loop()``. """ def __init__( self, name: str, *, scene_cam_name: str = "scene_camera", teleop_arm: Optional[str] = None, # None / "right" / "left" rerun_port: int = 9092, resolution: str = "HD1080", jpeg_target_w: int = 960, jpeg_quality: int = 75, exit_on_yellow: bool = True, ): self.name = name self.scene_cam_name = scene_cam_name self.teleop_arm = teleop_arm self.rerun_port = rerun_port self.resolution = resolution self.jpeg_target_w = jpeg_target_w self.jpeg_quality = jpeg_quality self.exit_on_yellow = exit_on_yellow self._cam = None self._rc = None self._follower = None self._frame_idx = 0 self._exit_requested = False self.K = None self.W = self.H = 0 # ── lifecycle ──────────────────────────────────────────────────────── def __enter__(self): from raiden.calibration.exo_calibrate import _open_zed_native serial = int(json.load(open(CAM_CFG_PATH))[self.scene_cam_name]["serial"]) self._cam, K, _dist, W, H = _open_zed_native(serial, self.resolution) self.K = np.asarray(K, dtype=np.float64) self.W, self.H = int(W), int(H) if self.teleop_arm is not None: self._rc = start_teleop(self.teleop_arm) self._follower = (self._rc.follower_r if self.teleop_arm == "right" else self._rc.follower_l) init_web(self.name, port=self.rerun_port) return self def __exit__(self, exc_type, exc, tb): if self._rc is not None: stop_teleop(self._rc) if self._cam is not None: self._cam.close() # ── properties ─────────────────────────────────────────────────────── @property def WH(self) -> tuple[int, int]: return self.W, self.H # ── camera ─────────────────────────────────────────────────────────── def grab(self) -> Optional[np.ndarray]: """Latest scene BGR frame, or ``None`` if grab failed. ZED's internal buffer holds only the freshest frame; a single grab returns it. A drain loop is *slower*, not faster — each grab() blocks waiting for the next frame. """ import pyzed.sl as sl if self._cam.grab(sl.RuntimeParameters()) != sl.ERROR_CODE.SUCCESS: return None image = sl.Mat() self._cam.retrieve_image(image, sl.VIEW.LEFT) return cv2.cvtColor(image.get_data(), cv2.COLOR_BGRA2BGR) # ── teleop ─────────────────────────────────────────────────────────── def follower_joints7(self) -> np.ndarray: if self._follower is None: raise RuntimeError("no teleop arm; pass teleop_arm to enable") return np.asarray(self._follower.get_joint_pos(), dtype=np.float64) def yellow_pressed(self) -> bool: """``True`` if the yellow leader top button was just pressed. Edge-triggered: returns True once per physical press. If ``exit_on_yellow`` is True (the default), the main loop polls this and exits; otherwise the callback owns polling and can treat yellow as a toggle. """ return self._rc is not None and self._rc.check_button_press() is not None def white_pressed(self) -> bool: """``True`` if the white leader bottom button was just pressed. Edge-triggered. Uses ``check_failure_button`` — the same path ``yam_control/test_record_joint_seq.py`` uses for WHITE-button joint capture, so it's the canonical "is the bottom button pressed right now" API on this hardware. """ if self._rc is None or not hasattr(self._rc, "check_failure_button"): return False try: return bool(self._rc.check_failure_button()) except Exception: return False def request_exit(self) -> None: """Signal :meth:`loop` to exit cleanly. For callback-driven exits when ``exit_on_yellow=False`` (e.g. record.py state machines).""" self._exit_requested = True # ── calibration ────────────────────────────────────────────────────── def lock_exo_pose(self, min_detections: int = 10, max_wait_s: float = 20.0) -> np.ndarray: """Pool ``min_detections`` exo-aruco frames, average → cam-in-rbase.""" print(f"locking cam pose: pooling {min_detections} aruco frames…") T = pool_exo_aruco( lambda: (self.grab(), self.K), min_detections=min_detections, max_wait_s=max_wait_s, ) print(f"✓ locked. T_cam_in_right_arm_base t = {T[:3, 3]}") return T # ── logging ────────────────────────────────────────────────────────── def log(self, path: str, img_bgr: np.ndarray) -> None: """JPEG-encode + ``rr.log`` at the current frame's timeline step. ``rr.set_time`` is called by :meth:`loop` once per iteration, so all ``log()`` calls inside the callback share the same step. """ rr.log(path, to_jpeg(img_bgr, target_w=self.jpeg_target_w, quality=self.jpeg_quality)) # ── main loop ──────────────────────────────────────────────────────── def loop(self, callback: Callable[[int, np.ndarray], None]) -> None: """Run the background grab+callback loop; main thread polls for exit. Callback signature: ``callback(frame_idx: int, bgr: np.ndarray)``. Exits on yellow-button press (teleop only) or ``KeyboardInterrupt``. """ stop = threading.Event() def _bg(): while not stop.is_set(): bgr = self.grab() if bgr is None: time.sleep(0.01); continue rr.set_time("step", sequence=self._frame_idx) callback(self._frame_idx, bgr) self._frame_idx += 1 th = threading.Thread(target=_bg, daemon=True) th.start() try: while True: if self._exit_requested: print("[exit requested]") break if self.exit_on_yellow and self.yellow_pressed(): print("[yellow] END") break time.sleep(0.05) finally: stop.set() th.join(timeout=1.0)