Add CSV plotting and episode recording functionality
- Implemented `plot_data_csv.py` to read CSV files and generate plots for velocity and linear acceleration signals. - Created `record1.py` for recording ESP32 IMU data, DIN/DOUT states, and RealSense camera images into episodes. - Enhanced `ProviderWorldIMUVelocityEstimator` to include stationary detection logic, resetting velocity when stationary. - Updated `EpisodeWriter` to save episode data with timestamped filenames for better organization.
This commit is contained in:
Brunsmeier
975
record1.py
Normal file
975
record1.py
Normal file
@ -0,0 +1,975 @@
|
||||
#!/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
|
||||
Reference in New Issue
Block a user