975 lines
27 KiB
Python
975 lines
27 KiB
Python
#!/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
|
|
|
|
# ---------------- 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 record_episode.py --port /dev/ttyUSB0 --baud 115200 --record-din 0 --dataset dataset --no-camera |