fd286e6a9e
- Introduced record_episode.py for capturing ESP32 IMU data and RealSense images. - Added SwitchChangeDetector and SwitchLevelController for managing recording triggers. - Enhanced ESP32Bridge with new methods for reading samples and latest packets. - Updated verify.py and visualise.py to improve serial communication stability. - Modified .gitignore to include dataset, csv, and png directories.
306 lines
8.1 KiB
Python
306 lines
8.1 KiB
Python
import struct
|
|
import time
|
|
import serial
|
|
|
|
|
|
|
|
SYNC = b"\xA5\x5A"
|
|
PACKET_FORMAT = "<2sBBBI10hB"
|
|
PACKET_SIZE = struct.calcsize(PACKET_FORMAT)
|
|
|
|
|
|
def decode_packet(packet):
|
|
fields = struct.unpack(PACKET_FORMAT, packet)
|
|
_, version, din_mask, dout_mask, time_ms, *rest = fields
|
|
raw = rest[:-1]
|
|
|
|
acc_scale = 16.0 / 32768.0
|
|
gyro_scale = 2000.0 / 32768.0
|
|
angle_scale = 180.0 / 32768.0
|
|
|
|
return {
|
|
"version": version,
|
|
"t_ms": time_ms,
|
|
"din": [(din_mask >> bit) & 1 for bit in range(4)],
|
|
"dout": [(dout_mask >> bit) & 1 for bit in range(2)],
|
|
"acc_g": [round(value * acc_scale, 5) for value in raw[0:3]],
|
|
"gyro_dps": [round(value * gyro_scale, 4) for value in raw[3:6]],
|
|
"angle_deg": [round(value * angle_scale, 4) for value in raw[6:9]],
|
|
"temp_c": round(raw[9] / 100.0, 2),
|
|
}
|
|
|
|
|
|
class ESP32Bridge:
|
|
def __init__(
|
|
self,
|
|
port,
|
|
baud=115200,
|
|
timeout=0.05,
|
|
startup_wait=4.0,
|
|
safe_stop_on_close=False,
|
|
reset_input_on_open=True,
|
|
reset_output_on_open=True,
|
|
):
|
|
self.port = port
|
|
self.baud = baud
|
|
self.timeout = timeout
|
|
self.startup_wait = startup_wait
|
|
self.safe_stop_on_close = safe_stop_on_close
|
|
self.reset_input_on_open = reset_input_on_open
|
|
self.reset_output_on_open = reset_output_on_open
|
|
|
|
self.ser = None
|
|
self.buffer = bytearray()
|
|
self.last_t_ms = None
|
|
|
|
def is_open(self):
|
|
return self.ser is not None and self.ser.is_open
|
|
|
|
def _require_open(self):
|
|
if not self.is_open():
|
|
raise RuntimeError("Serial port is not open.")
|
|
|
|
def open(self):
|
|
if self.is_open():
|
|
return self
|
|
|
|
self.ser = serial.Serial()
|
|
self.ser.port = self.port
|
|
self.ser.baudrate = self.baud
|
|
self.ser.timeout = self.timeout
|
|
self.ser.write_timeout = self.timeout
|
|
|
|
# Avoid ESP32 auto reset / boot mode weirdness.
|
|
self.ser.dtr = False
|
|
self.ser.rts = False
|
|
self.ser.open()
|
|
|
|
self.ser.dtr = False
|
|
self.ser.rts = False
|
|
|
|
if self.reset_input_on_open:
|
|
self.ser.reset_input_buffer()
|
|
|
|
if self.reset_output_on_open:
|
|
self.ser.reset_output_buffer()
|
|
|
|
self.buffer.clear()
|
|
self.last_t_ms = None
|
|
|
|
if self.startup_wait > 0:
|
|
time.sleep(self.startup_wait)
|
|
|
|
return self
|
|
|
|
def close(self):
|
|
if self.ser is None:
|
|
return
|
|
|
|
if self.safe_stop_on_close and self.ser.is_open:
|
|
try:
|
|
self.stop_gripper()
|
|
time.sleep(0.1)
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
if self.ser.is_open:
|
|
self.ser.close()
|
|
finally:
|
|
self.ser = None
|
|
|
|
def __enter__(self):
|
|
return self.open()
|
|
|
|
def __exit__(self, exc_type, exc, tb):
|
|
self.close()
|
|
|
|
def send_gripper_command(self, mode):
|
|
self._require_open()
|
|
|
|
if mode == 1:
|
|
self.ser.write(b"o")
|
|
elif mode == 2:
|
|
self.ser.write(b"c")
|
|
else:
|
|
self.ser.write(b"s")
|
|
|
|
self.ser.flush()
|
|
|
|
def open_gripper(self):
|
|
self.send_gripper_command(1)
|
|
|
|
def close_gripper(self):
|
|
self.send_gripper_command(2)
|
|
|
|
def stop_gripper(self):
|
|
self.send_gripper_command(0)
|
|
|
|
def read_packet(self):
|
|
"""
|
|
Return one decoded packet dict.
|
|
This blocks until a valid packet is found.
|
|
"""
|
|
self._require_open()
|
|
|
|
while True:
|
|
chunk = self.ser.read(self.ser.in_waiting or 1)
|
|
if chunk:
|
|
self.buffer.extend(chunk)
|
|
|
|
packet = self._try_pop_packet()
|
|
if packet is not None:
|
|
return decode_packet(packet)
|
|
|
|
def _try_pop_packet(self):
|
|
while len(self.buffer) >= PACKET_SIZE:
|
|
sync_index = self.buffer.find(SYNC)
|
|
|
|
if sync_index < 0:
|
|
# Keep last byte in case it is the first half of SYNC.
|
|
del self.buffer[:-1]
|
|
return None
|
|
|
|
if sync_index:
|
|
del self.buffer[:sync_index]
|
|
|
|
if len(self.buffer) < PACKET_SIZE:
|
|
return None
|
|
|
|
packet = bytes(self.buffer[:PACKET_SIZE])
|
|
|
|
if (sum(packet[:-1]) & 0xFF) != packet[-1]:
|
|
del self.buffer[0]
|
|
continue
|
|
|
|
del self.buffer[:PACKET_SIZE]
|
|
return packet
|
|
|
|
return None
|
|
|
|
def read_packet_with_frequency(self):
|
|
data = self.read_packet()
|
|
|
|
if self.last_t_ms is None:
|
|
freq_hz = 0.0
|
|
else:
|
|
delta = (data["t_ms"] - self.last_t_ms) & 0xFFFFFFFF
|
|
freq_hz = 1000.0 / delta if delta else 0.0
|
|
|
|
self.last_t_ms = data["t_ms"]
|
|
data["freq_hz"] = freq_hz
|
|
return data
|
|
|
|
def iter_packets(self, seconds=None, max_samples=None, with_frequency=True):
|
|
start = time.monotonic()
|
|
count = 0
|
|
|
|
while True:
|
|
if seconds is not None and time.monotonic() - start >= seconds:
|
|
break
|
|
|
|
if max_samples is not None and count >= max_samples:
|
|
break
|
|
|
|
if with_frequency:
|
|
yield self.read_packet_with_frequency()
|
|
else:
|
|
yield self.read_packet()
|
|
|
|
count += 1
|
|
def read_samples(self, seconds=None, max_samples=None, max_gap_ms=1000):
|
|
"""
|
|
Read multiple packets and return flattened samples.
|
|
|
|
Suitable for visualise.py / verify.py.
|
|
"""
|
|
self._require_open()
|
|
|
|
samples = []
|
|
bad_delta_count = 0
|
|
elapsed_s = 0.0
|
|
start = time.monotonic()
|
|
|
|
while True:
|
|
if seconds is not None and time.monotonic() - start >= seconds:
|
|
break
|
|
|
|
if max_samples is not None and len(samples) >= max_samples:
|
|
break
|
|
|
|
data = self.read_packet()
|
|
|
|
if self.last_t_ms is None:
|
|
freq_hz = 0.0
|
|
else:
|
|
delta = (data["t_ms"] - self.last_t_ms) & 0xFFFFFFFF
|
|
if 0 < delta <= max_gap_ms:
|
|
elapsed_s += delta / 1000.0
|
|
freq_hz = 1000.0 / delta
|
|
else:
|
|
bad_delta_count += 1
|
|
freq_hz = 0.0
|
|
|
|
self.last_t_ms = data["t_ms"]
|
|
|
|
sample = flatten_packet(data, elapsed_s, freq_hz)
|
|
samples.append(sample)
|
|
|
|
return samples, bad_delta_count
|
|
|
|
def read_latest_packet(self, max_drain=200):
|
|
"""
|
|
Return the newest decoded packet currently available.
|
|
|
|
It drains old buffered packets so the PC side follows the latest DIN/DOUT state.
|
|
Useful when camera/image saving is slower than ESP32 50Hz output.
|
|
"""
|
|
self._require_open()
|
|
|
|
latest = self.read_packet()
|
|
drained = 0
|
|
|
|
while drained < max_drain:
|
|
# Pull all currently available bytes from OS serial buffer.
|
|
waiting = self.ser.in_waiting
|
|
if waiting <= 0 and len(self.buffer) < PACKET_SIZE:
|
|
break
|
|
|
|
if waiting > 0:
|
|
chunk = self.ser.read(waiting)
|
|
if chunk:
|
|
self.buffer.extend(chunk)
|
|
|
|
packet = self._try_pop_packet()
|
|
|
|
if packet is None:
|
|
break
|
|
|
|
latest = decode_packet(packet)
|
|
drained += 1
|
|
|
|
return latest, drained
|
|
|
|
def flatten_packet(data, t_s, freq_hz):
|
|
return {
|
|
"t_s": t_s,
|
|
"acc_x_g": data["acc_g"][0],
|
|
"acc_y_g": data["acc_g"][1],
|
|
"acc_z_g": data["acc_g"][2],
|
|
"gyro_x_dps": data["gyro_dps"][0],
|
|
"gyro_y_dps": data["gyro_dps"][1],
|
|
"gyro_z_dps": data["gyro_dps"][2],
|
|
"angle_x_deg": data["angle_deg"][0],
|
|
"angle_y_deg": data["angle_deg"][1],
|
|
"angle_z_deg": data["angle_deg"][2],
|
|
"temp_c": data["temp_c"],
|
|
"freq_hz": freq_hz,
|
|
"din0": data["din"][0],
|
|
"din1": data["din"][1],
|
|
"din2": data["din"][2],
|
|
"din3": data["din"][3],
|
|
"dout0": data["dout"][0],
|
|
"dout1": data["dout"][1],
|
|
}
|
|
|