#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ record_episode.py PC-side episode recorder. Records: - ESP32 IMU data: acc / gyro / angle / temperature - ESP32 DIN state: GPIO25 / GPIO26 / GPIO27 / GPIO32 - ESP32 DOUT logical state: GPIO18 / GPIO19 relay command state - inferred scissor state: open / close / stop - RealSense D405 RGB image - RealSense D405 depth image - Filtered IMU integrated velocity Default DIN mapping from ESP32 main.py: DIN_PINS = (25, 26, 27, 32) Therefore: din[0] = GPIO25 din[1] = GPIO26 din[2] = GPIO27 # already used as physical scissor switch din[3] = GPIO32 Default recording trigger: din[0] / GPIO25 Usage: Test ESP32 switch + CSV only: python3 record_episode.py --port /dev/ttyUSB0 --no-camera Record with D405 RGB + depth: python3 record_episode.py --port /dev/ttyUSB0 Use GPIO26 as record trigger: python3 record_episode.py --port /dev/ttyUSB0 --record-din 1 """ import argparse import csv import time import math from pathlib import Path from datetime import datetime import cv2 import numpy as np try: import pyrealsense2 as rs except ImportError: rs = None from esp_bridge import ESP32Bridge # ============================================================ # Switch edge detector # ============================================================ class SwitchChangeDetector: """ Detect any valid state change from a toggle switch. For a latching/toggle switch: 1 -> 0 triggers once 0 -> 1 triggers once This is different from a push button, where only "pressed" edge is used. """ def __init__(self, debounce_s=0.25): self.debounce_s = debounce_s self.last_value = None self.last_edge_time = 0.0 def update(self, raw_value): now = time.monotonic() # First sample only initializes state. # Do not trigger recording immediately when program starts. if self.last_value is None: self.last_value = raw_value return False if raw_value == self.last_value: return False if now - self.last_edge_time < self.debounce_s: return False old_value = self.last_value self.last_value = raw_value self.last_edge_time = now print(f"[SWITCH CHANGE] {old_value} -> {raw_value}") return True # ============================================================ # Switch level controller # ============================================================ class SwitchLevelController: """ Level-control recording switch. raw_value == active_value -> should record raw_value != active_value -> should stop For your current requirement: DIN0 = 1 -> REC DIN0 = 0 -> IDLE """ def __init__(self, active_value=1, debounce_s=0.25): self.active_value = active_value self.debounce_s = debounce_s self.last_raw_value = None self.stable_value = None self.last_change_time = 0.0 def update(self, raw_value): now = time.monotonic() # First sample initializes stable state. # If program starts while switch is already 1, it will start recording. if self.last_raw_value is None: self.last_raw_value = raw_value self.stable_value = raw_value return raw_value == self.active_value # Raw value changed. Start debounce timer. if raw_value != self.last_raw_value: self.last_raw_value = raw_value self.last_change_time = now return None # Raw value has remained the same. Accept it after debounce time. if raw_value != self.stable_value: if now - self.last_change_time >= self.debounce_s: old_value = self.stable_value self.stable_value = raw_value print(f"[SWITCH LEVEL] {old_value} -> {self.stable_value}") return self.stable_value == self.active_value return None # ============================================================ # RealSense D405 camera # ============================================================ class D405Camera: def __init__( self, width=640, height=480, fps=30, enable_depth=True, warmup_frames=15, ): if rs is None: raise RuntimeError( "pyrealsense2 is not installed. " "Install it or run with --no-camera." ) self.width = width self.height = height self.fps = fps self.enable_depth = enable_depth self.warmup_frames = warmup_frames self.pipeline = rs.pipeline() self.config = rs.config() self.align = None def open(self): self.config.enable_stream( rs.stream.color, self.width, self.height, rs.format.bgr8, self.fps, ) if self.enable_depth: self.config.enable_stream( rs.stream.depth, self.width, self.height, rs.format.z16, self.fps, ) self.align = rs.align(rs.stream.color) self.pipeline.start(self.config) # Warm up auto exposure. for _ in range(self.warmup_frames): self.pipeline.wait_for_frames() return self def read(self): frames = self.pipeline.wait_for_frames() if self.enable_depth and self.align is not None: frames = self.align.process(frames) color_frame = frames.get_color_frame() if not color_frame: return None, None color = np.asanyarray(color_frame.get_data()) depth = None if self.enable_depth: depth_frame = frames.get_depth_frame() if depth_frame: depth = np.asanyarray(depth_frame.get_data()) return color, depth def close(self): try: self.pipeline.stop() except Exception: pass def __enter__(self): return self.open() def __exit__(self, exc_type, exc, tb): self.close() # ============================================================ # Simple IMU velocity estimator Updated # ============================================================ class ProviderWorldIMUVelocityEstimator: """ Provider-style gravity/support compensation + world-frame velocity integration. Pipeline: 1. raw acc_body_g - gravity_body_g -> linear_acc_body_g 2. estimate residual body-frame acceleration bias 3. linear_acc_body_ms2 -> linear_acc_world_ms2 using R_body_to_world 4. integrate velocity in world frame Output: velocity_ms: world-frame velocity, unit m/s linear_acc_ms2: world-frame linear acceleration, unit m/s^2 acc_bias_ms2: body-frame residual acceleration bias, unit m/s^2 """ def __init__( self, gravity_sign=1.0, bias_seconds=1.0, acc_deadband=0.10, ): self.gravity_sign = gravity_sign self.bias_seconds = bias_seconds self.acc_deadband = acc_deadband self.start_t_ms = None self.last_t_ms = None self.bias_ready = False self.bias_samples = [] self.acc_bias_body_ms2 = np.zeros(3, dtype=float) self.v_world = np.zeros(3, dtype=float) self.last_acc_world_ms2 = np.zeros(3, dtype=float) def reset(self): self.start_t_ms = None self.last_t_ms = None self.bias_ready = False self.bias_samples = [] self.acc_bias_body_ms2[:] = 0.0 self.v_world[:] = 0.0 self.last_acc_world_ms2[:] = 0.0 def gravity_body_g(self, roll_deg, pitch_deg): roll = math.radians(roll_deg) pitch = math.radians(pitch_deg) g_x = -math.sin(pitch) g_y = math.cos(pitch) * math.sin(roll) g_z = math.cos(pitch) * math.cos(roll) return np.array([g_x, g_y, g_z], dtype=float) * self.gravity_sign def rot_body_to_world(self, roll_deg, pitch_deg, yaw_deg): roll = math.radians(roll_deg) pitch = math.radians(pitch_deg) yaw = math.radians(yaw_deg) cr, sr = math.cos(roll), math.sin(roll) cp, sp = math.cos(pitch), math.sin(pitch) cy, sy = math.cos(yaw), math.sin(yaw) # R_body_to_world = Rz(yaw) @ Ry(pitch) @ Rx(roll) return np.array([ [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr], [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr], [-sp, cp * sr, cp * cr], ], dtype=float) def update(self, packet): t_ms = packet["t_ms"] if self.start_t_ms is None: self.start_t_ms = t_ms acc_body_g = np.array(packet["acc_g"], dtype=float) roll_deg = packet["angle_deg"][0] pitch_deg = packet["angle_deg"][1] yaw_deg = packet["angle_deg"][2] # ---------------- gravity/support removal in body frame ---------------- gravity_g = self.gravity_body_g( roll_deg=roll_deg, pitch_deg=pitch_deg, ) linear_acc_body_g_raw = acc_body_g - gravity_g linear_acc_body_ms2_raw = linear_acc_body_g_raw * 9.80665 elapsed_s = ((t_ms - self.start_t_ms) & 0xFFFFFFFF) / 1000.0 # ---------------- bias estimation in body frame ---------------- # During this period, velocity stays zero. if not self.bias_ready: self.bias_samples.append(linear_acc_body_ms2_raw.copy()) if elapsed_s >= self.bias_seconds: if self.bias_samples: self.acc_bias_body_ms2 = np.mean( np.array(self.bias_samples, dtype=float), axis=0, ) self.bias_ready = True print( "[IMU] Provider body-frame acc bias estimated: " f"{self.acc_bias_body_ms2}" ) self.last_t_ms = t_ms self.last_acc_world_ms2[:] = 0.0 return { "gravity_g": gravity_g.copy(), "linear_acc_ms2": self.last_acc_world_ms2.copy(), "velocity_ms": self.v_world.copy(), "acc_bias_ms2": self.acc_bias_body_ms2.copy(), "bias_ready": self.bias_ready, } # ---------------- remove body-frame residual bias ---------------- linear_acc_body_ms2 = linear_acc_body_ms2_raw - self.acc_bias_body_ms2 # linear_acc_body_ms2 = linear_acc_body_ms2_raw # ---------------- deadband before rotation ---------------- # Norm is unchanged by pure rotation, so before/after rotation both okay. # if np.linalg.norm(linear_acc_body_ms2) < self.acc_deadband: # linear_acc_body_ms2[:] = 0.0 # ---------------- body frame -> world frame ---------------- R_bw = self.rot_body_to_world( roll_deg=roll_deg, pitch_deg=pitch_deg, yaw_deg=yaw_deg, ) linear_acc_world_ms2 = R_bw @ linear_acc_body_ms2 # ---------------- dt ---------------- delta_ms = (t_ms - self.last_t_ms) & 0xFFFFFFFF self.last_t_ms = t_ms if 0 < delta_ms <= 200: dt = delta_ms / 1000.0 self.v_world += linear_acc_world_ms2 * dt self.last_acc_world_ms2 = linear_acc_world_ms2.copy() return { "gravity_g": gravity_g.copy(), "linear_acc_ms2": self.last_acc_world_ms2.copy(), "velocity_ms": self.v_world.copy(), "acc_bias_ms2": self.acc_bias_body_ms2.copy(), "bias_ready": self.bias_ready, } # ============================================================ # Scissor state # ============================================================ def infer_scissor_state(packet): """ DOUT logical mapping: dout[0] = GPIO18 = relay IN1 dout[1] = GPIO19 = relay IN2 Hardware behavior: IN1 active -> CLOSE IN2 active -> OPEN Important: packet["dout"] comes from ESP32 dout_shadow. It is a logical active state, not the raw GPIO voltage. """ dout0 = packet["dout"][0] dout1 = packet["dout"][1] if dout0 == 1 and dout1 == 0: return "close" if dout0 == 0 and dout1 == 1: return "open" if dout0 == 0 and dout1 == 0: return "stop" return "invalid_both_active" # ============================================================ # Episode writer # ============================================================ class EpisodeWriter: def __init__( self, root_dir="dataset", save_rgb=True, save_depth=True, jpeg_quality=95, flush_every=10, ): self.root_dir = Path(root_dir) self.save_rgb = save_rgb self.save_depth = save_depth self.jpeg_quality = jpeg_quality self.flush_every = flush_every self.root_dir.mkdir(parents=True, exist_ok=True) self.recording = False self.episode_dir = None self.rgb_dir = None self.depth_dir = None self.csv_file = None self.writer = None self.frame_id = 0 def start(self): if self.recording: return episode_name = "episode_" + datetime.now().strftime("%Y%m%d_%H%M%S") self.episode_dir = self.root_dir / episode_name self.rgb_dir = self.episode_dir / "rgb" self.depth_dir = self.episode_dir / "depth" self.episode_dir.mkdir(parents=True, exist_ok=True) if self.save_rgb: self.rgb_dir.mkdir(parents=True, exist_ok=True) if self.save_depth: self.depth_dir.mkdir(parents=True, exist_ok=True) csv_path = self.episode_dir / "data.csv" self.csv_file = open(csv_path, "w", newline="", encoding="utf-8") fieldnames = [ "frame_id", "t_wall", "t_monotonic", "t_esp_ms", "din0_gpio25_record_switch", "din1_gpio26", "din2_gpio27_scissor_switch", "din3_gpio32", "dout0_gpio18_in1", "dout1_gpio19_in2", "scissor_state", "acc_x_g", "acc_y_g", "acc_z_g", "gyro_x_dps", "gyro_y_dps", "gyro_z_dps", "angle_x_deg", "angle_y_deg", "angle_z_deg", "temp_c", "vel_x_ms", "vel_y_ms", "vel_z_ms", "lin_acc_x_ms2", "lin_acc_y_ms2", "lin_acc_z_ms2", "acc_bias_x_ms2", "acc_bias_y_ms2", "acc_bias_z_ms2", "bias_ready", "rgb_path", "depth_path", ] self.writer = csv.DictWriter(self.csv_file, fieldnames=fieldnames) self.writer.writeheader() self.frame_id = 0 self.recording = True print(f"[REC START] {self.episode_dir}") def stop(self): if not self.recording: return self.recording = False if self.csv_file is not None: self.csv_file.flush() self.csv_file.close() print(f"[REC STOP] {self.episode_dir}") self.csv_file = None self.writer = None self.episode_dir = None self.rgb_dir = None self.depth_dir = None self.frame_id = 0 def toggle(self): if self.recording: self.stop() else: self.start() def write(self, packet, imu_state, color_image=None, depth_image=None): if not self.recording: return rgb_rel = "" depth_rel = "" if self.save_rgb and color_image is not None: rgb_name = f"{self.frame_id:06d}.jpg" rgb_path = self.rgb_dir / rgb_name ok = cv2.imwrite( str(rgb_path), color_image, [int(cv2.IMWRITE_JPEG_QUALITY), int(self.jpeg_quality)], ) if ok: rgb_rel = "rgb/" + rgb_name if self.save_depth and depth_image is not None: depth_name = f"{self.frame_id:06d}.png" depth_path = self.depth_dir / depth_name ok = cv2.imwrite(str(depth_path), depth_image) if ok: depth_rel = "depth/" + depth_name row = { "frame_id": self.frame_id, "t_wall": time.time(), "t_monotonic": time.monotonic(), "t_esp_ms": packet["t_ms"], "din0_gpio25_record_switch": packet["din"][0], "din1_gpio26": packet["din"][1], "din2_gpio27_scissor_switch": packet["din"][2], "din3_gpio32": packet["din"][3], "dout0_gpio18_in1": packet["dout"][0], "dout1_gpio19_in2": packet["dout"][1], "scissor_state": infer_scissor_state(packet), "acc_x_g": packet["acc_g"][0], "acc_y_g": packet["acc_g"][1], "acc_z_g": packet["acc_g"][2], "gyro_x_dps": packet["gyro_dps"][0], "gyro_y_dps": packet["gyro_dps"][1], "gyro_z_dps": packet["gyro_dps"][2], "angle_x_deg": packet["angle_deg"][0], "angle_y_deg": packet["angle_deg"][1], "angle_z_deg": packet["angle_deg"][2], "temp_c": packet["temp_c"], "vel_x_ms": float(imu_state["velocity_ms"][0]), "vel_y_ms": float(imu_state["velocity_ms"][1]), "vel_z_ms": float(imu_state["velocity_ms"][2]), "lin_acc_x_ms2": float(imu_state["linear_acc_ms2"][0]), "lin_acc_y_ms2": float(imu_state["linear_acc_ms2"][1]), "lin_acc_z_ms2": float(imu_state["linear_acc_ms2"][2]), "acc_bias_x_ms2": float(imu_state["acc_bias_ms2"][0]), "acc_bias_y_ms2": float(imu_state["acc_bias_ms2"][1]), "acc_bias_z_ms2": float(imu_state["acc_bias_ms2"][2]), "bias_ready": int(imu_state["bias_ready"]), "rgb_path": rgb_rel, "depth_path": depth_rel, } self.writer.writerow(row) self.frame_id += 1 if self.flush_every > 0 and self.frame_id % self.flush_every == 0: self.csv_file.flush() # ============================================================ # Main recorder # ============================================================ def run(args): if args.record_din == 2: print( "[WARNING] You selected din[2] / GPIO27 as recording trigger. " "But GPIO27 is already used as physical scissor switch." ) # record_switch = SwitchChangeDetector( # debounce_s=args.debounce, # ) #in case we need this former switch change detector, we can keep it here for reference. But the level controller is more suitable for this use case. record_switch = SwitchLevelController( active_value=args.record_active_value, debounce_s=args.debounce, ) # velocity_estimator = ProviderIMUVelocityEstimator( # gravity_sign=args.gravity_sign, # bias_seconds=args.bias_seconds, # acc_deadband=args.acc_deadband, # vel_decay=args.velocity_decay, # ) velocity_estimator = ProviderWorldIMUVelocityEstimator( gravity_sign=args.gravity_sign, bias_seconds=args.bias_seconds, acc_deadband=args.acc_deadband, ) writer = EpisodeWriter( root_dir=args.dataset, save_rgb=not args.no_rgb, save_depth=not args.no_depth, jpeg_quality=args.jpeg_quality, flush_every=args.flush_every, ) camera = None try: if args.no_camera: print("[CAMERA] disabled by --no-camera.") else: camera = D405Camera( width=args.width, height=args.height, fps=args.fps, enable_depth=not args.no_depth, warmup_frames=args.camera_warmup, ) camera.open() print("[CAMERA] D405 opened.") with ESP32Bridge( port=args.port, baud=args.baud, startup_wait=args.startup_wait, safe_stop_on_close=False, ) as bridge: print(f"[ESP32] Reading {args.port} at {args.baud} baud.") print( "[SWITCH] record trigger: " f"din[{args.record_din}], " f"mode=level-control, " f"active_value={args.record_active_value}, " f"debounce={args.debounce}s" ) print( f"[INFO] DIN[{args.record_din}]={args.record_active_value} -> REC, " f"DIN[{args.record_din}]!={args.record_active_value} -> IDLE." ) print("[INFO] Ctrl+C to exit.") packet_count = 0 last_status_t = time.monotonic() while True: packet, drained = bridge.read_latest_packet() packet_count += 1 + drained serial_waiting = bridge.ser.in_waiting if bridge.ser is not None else -1 raw_record_value = packet["din"][args.record_din] desired_recording = record_switch.update(raw_record_value) if desired_recording is not None: if desired_recording and not writer.recording: writer.start() velocity_estimator.reset() elif not desired_recording and writer.recording: writer.stop() imu_state = velocity_estimator.update(packet) color_image = None depth_image = None if writer.recording and camera is not None: color_image, depth_image = camera.read() writer.write( packet=packet, imu_state=imu_state, color_image=color_image, depth_image=depth_image, ) now = time.monotonic() if now - last_status_t >= args.status_period: last_status_t = now rec_state = "REC" if writer.recording else "IDLE" print( f"[{rec_state}] " f"packets={packet_count} " f"DIN={packet['din']} " f"DOUT={packet['dout']} " f"scissor={infer_scissor_state(packet)} " f"frames={writer.frame_id} " f"drained={drained} " f"serial_waiting={serial_waiting}" ) except KeyboardInterrupt: print("\n[EXIT] Interrupted by user.") finally: writer.stop() if camera is not None: camera.close() print("[CAMERA] closed.") def parse_args(): parser = argparse.ArgumentParser( description="Record ESP32 IMU/DIN/DOUT + D405 RGB/depth into episodes." ) parser.add_argument( "--port", required=True, help="ESP32 serial port, for example /dev/ttyUSB0", ) parser.add_argument( "--baud", type=int, default=115200, help="ESP32 USB serial baudrate.", ) parser.add_argument( "--dataset", default="dataset", help="Output dataset root directory.", ) parser.add_argument( "--record-din", type=int, default=0, choices=[0, 1, 2, 3], help=( "DIN index used as record start/stop switch. " "DIN_PINS=(25,26,27,32). " "0=GPIO25, 1=GPIO26, 2=GPIO27, 3=GPIO32. " "GPIO27/din2 is already used as scissor physical switch, " "so default is din0/GPIO25." ), ) parser.add_argument( "--record-active-value", type=int, default=1, choices=[0, 1], help="Record when selected DIN equals this value. Default: 1, meaning DIN=1 records and DIN=0 stops.", ) parser.add_argument( "--active-low", dest="active_low", action="store_true", default=True, help="Record switch is low-level active. Default: True.", ) parser.add_argument( "--active-high", dest="active_low", action="store_false", help="Use high-level active record switch.", ) parser.add_argument( "--debounce", type=float, default=0.25, help="Record switch debounce time in seconds.", ) parser.add_argument( "--startup-wait", type=float, default=4.0, help="Wait time after opening ESP32 serial port.", ) parser.add_argument( "--no-camera", action="store_true", help="Disable RealSense camera. Useful for testing DIN toggle and CSV only.", ) parser.add_argument( "--width", type=int, default=640, help="RealSense color/depth width.", ) parser.add_argument( "--height", type=int, default=480, help="RealSense color/depth height.", ) parser.add_argument( "--fps", type=int, default=30, help="RealSense FPS.", ) parser.add_argument( "--camera-warmup", type=int, default=15, help="Frames to discard after starting camera.", ) parser.add_argument( "--no-rgb", action="store_true", help="Do not save RGB images.", ) parser.add_argument( "--no-depth", action="store_true", help="Do not save depth images.", ) parser.add_argument( "--jpeg-quality", type=int, default=95, help="JPEG quality for RGB image saving.", ) # parser.add_argument( # "--velocity-decay", # type=float, # default=0.995, # help="Decay factor for simple raw IMU velocity integration.", # ) parser.add_argument( "--flush-every", type=int, default=10, help="Flush CSV every N frames. Use 1 for safer but slower writing.", ) parser.add_argument( "--status-period", type=float, default=1.0, help="Print status every N seconds.", ) parser.add_argument( "--gravity-sign", type=float, default=1.0, choices=[1.0, -1.0], help="Use +1 if horizontal static acc_z is about +1g; use -1 if about -1g.", ) parser.add_argument( "--bias-seconds", type=float, default=1.0, help="Use the first N seconds after episode start to estimate residual acceleration bias.", ) parser.add_argument( "--acc-deadband", type=float, default=0.10, help="Set linear acceleration to zero if norm is below this value, unit m/s^2.", ) return parser.parse_args() if __name__ == "__main__": run(parse_args()) # python3 record1.py --port /dev/ttyUSB0 --baud 115200 --record-din 0 --dataset dataset --no-camera