Update to high speed encoder stream
This commit is contained in:
283
Python/parse_binary_stream.py
Normal file
283
Python/parse_binary_stream.py
Normal file
@@ -0,0 +1,283 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import struct
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
|
||||
|
||||
BAUD_RATE = 921_600
|
||||
UINT32_MASK = 0xFFFFFFFF
|
||||
|
||||
TURN_MAGIC = b"\xE7\x54\xC3\xA1"
|
||||
PULSE16_MAGIC = b"\xE7\x50\xC3\xA1"
|
||||
PULSE32_MAGIC = b"\xE7\x70\xC3\xA1"
|
||||
GP_MAGIC = b"\xE7\x47\xC3\xA1"
|
||||
MAGICS = {
|
||||
TURN_MAGIC: "turn",
|
||||
PULSE16_MAGIC: "pulse16",
|
||||
PULSE32_MAGIC: "pulse32",
|
||||
GP_MAGIC: "gp",
|
||||
}
|
||||
|
||||
GP_COLUMNS = ("gp0_falling", "gp1_falling")
|
||||
CSV_COLUMNS = ("time_us", "turn", "pulse", *GP_COLUMNS)
|
||||
|
||||
|
||||
class EventTable:
|
||||
def __init__(self) -> None:
|
||||
self.offset: int | None = None
|
||||
self.rows: dict[int, list[int]] = {}
|
||||
|
||||
def add(self, timestamp: int, column: str) -> None:
|
||||
if self.offset is None:
|
||||
self.offset = timestamp
|
||||
|
||||
time_us = (timestamp - self.offset) & UINT32_MASK
|
||||
row = self.rows.setdefault(time_us, [0] * (len(CSV_COLUMNS) - 1))
|
||||
row[CSV_COLUMNS.index(column) - 1] = 1
|
||||
|
||||
def write_csv(self, path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", newline="") as file:
|
||||
writer = csv.writer(file)
|
||||
writer.writerow(CSV_COLUMNS)
|
||||
for time_us in sorted(self.rows):
|
||||
writer.writerow((time_us, *self.rows[time_us]))
|
||||
|
||||
|
||||
def choose_port() -> str | None:
|
||||
from serial.tools import list_ports
|
||||
|
||||
ports = list(list_ports.comports())
|
||||
if not ports:
|
||||
print("No serial ports found.")
|
||||
return None
|
||||
|
||||
print("Available serial ports:")
|
||||
for index, port in enumerate(ports, start=1):
|
||||
description = port.description or "serial port"
|
||||
print(f" {index}. {port.device} - {description}")
|
||||
|
||||
while True:
|
||||
choice = input("Choose a port number, or q to quit: ").strip().lower()
|
||||
if choice == "q":
|
||||
return None
|
||||
|
||||
try:
|
||||
index = int(choice)
|
||||
except ValueError:
|
||||
print("Please enter a port number.")
|
||||
continue
|
||||
|
||||
if 1 <= index <= len(ports):
|
||||
return ports[index - 1].device
|
||||
|
||||
print(f"Please choose a number between 1 and {len(ports)}.")
|
||||
|
||||
|
||||
def deadline_expired(deadline: float | None) -> bool:
|
||||
return deadline is not None and time.monotonic() >= deadline
|
||||
|
||||
|
||||
def read_exact(
|
||||
stream: BinaryIO,
|
||||
size: int,
|
||||
*,
|
||||
deadline: float | None,
|
||||
eof_on_empty: bool,
|
||||
) -> bytes | None:
|
||||
data = bytearray()
|
||||
while len(data) < size:
|
||||
chunk = stream.read(size - len(data))
|
||||
if not chunk:
|
||||
if eof_on_empty or deadline_expired(deadline):
|
||||
return None
|
||||
continue
|
||||
data.extend(chunk)
|
||||
return bytes(data)
|
||||
|
||||
|
||||
def read_magic(
|
||||
stream: BinaryIO,
|
||||
*,
|
||||
deadline: float | None,
|
||||
eof_on_empty: bool,
|
||||
) -> str | None:
|
||||
window = bytearray()
|
||||
|
||||
while True:
|
||||
byte = stream.read(1)
|
||||
if not byte:
|
||||
if eof_on_empty or deadline_expired(deadline):
|
||||
return None
|
||||
continue
|
||||
|
||||
window.extend(byte)
|
||||
if len(window) > 4:
|
||||
del window[0]
|
||||
|
||||
if len(window) == 4:
|
||||
frame_type = MAGICS.get(bytes(window))
|
||||
if frame_type is not None:
|
||||
return frame_type
|
||||
|
||||
|
||||
def parse_stream(
|
||||
stream: BinaryIO,
|
||||
events: EventTable,
|
||||
duration_s: float | None = None,
|
||||
eof_on_empty: bool = True,
|
||||
) -> int:
|
||||
previous_pulse_timestamp: int | None = None
|
||||
deadline = None if duration_s is None else time.monotonic() + duration_s
|
||||
frames = 0
|
||||
|
||||
while deadline is None or time.monotonic() < deadline:
|
||||
frame_type = read_magic(
|
||||
stream,
|
||||
deadline=deadline,
|
||||
eof_on_empty=eof_on_empty,
|
||||
)
|
||||
if frame_type is None:
|
||||
break
|
||||
|
||||
if frame_type == "turn":
|
||||
payload = read_exact(
|
||||
stream,
|
||||
4,
|
||||
deadline=deadline,
|
||||
eof_on_empty=eof_on_empty,
|
||||
)
|
||||
if payload is None:
|
||||
break
|
||||
timestamp = struct.unpack("<I", payload)[0]
|
||||
previous_pulse_timestamp = timestamp
|
||||
events.add(timestamp, "turn")
|
||||
elif frame_type == "pulse16":
|
||||
payload = read_exact(
|
||||
stream,
|
||||
16 * 2,
|
||||
deadline=deadline,
|
||||
eof_on_empty=eof_on_empty,
|
||||
)
|
||||
if payload is None:
|
||||
break
|
||||
if previous_pulse_timestamp is not None:
|
||||
for delta in struct.unpack("<16H", payload):
|
||||
if delta == 0:
|
||||
continue
|
||||
previous_pulse_timestamp = (
|
||||
previous_pulse_timestamp + delta
|
||||
) & UINT32_MASK
|
||||
events.add(previous_pulse_timestamp, "pulse")
|
||||
elif frame_type == "pulse32":
|
||||
payload = read_exact(
|
||||
stream,
|
||||
16 * 4,
|
||||
deadline=deadline,
|
||||
eof_on_empty=eof_on_empty,
|
||||
)
|
||||
if payload is None:
|
||||
break
|
||||
if previous_pulse_timestamp is not None:
|
||||
for delta in struct.unpack("<16I", payload):
|
||||
if delta == 0:
|
||||
continue
|
||||
previous_pulse_timestamp = (
|
||||
previous_pulse_timestamp + delta
|
||||
) & UINT32_MASK
|
||||
events.add(previous_pulse_timestamp, "pulse")
|
||||
elif frame_type == "gp":
|
||||
payload = read_exact(
|
||||
stream,
|
||||
5,
|
||||
deadline=deadline,
|
||||
eof_on_empty=eof_on_empty,
|
||||
)
|
||||
if payload is None:
|
||||
break
|
||||
channel = payload[0]
|
||||
timestamp = struct.unpack("<I", payload[1:])[0]
|
||||
if channel < len(GP_COLUMNS):
|
||||
events.add(timestamp, GP_COLUMNS[channel])
|
||||
|
||||
frames += 1
|
||||
|
||||
return frames
|
||||
|
||||
|
||||
def default_output_path() -> Path:
|
||||
started_at = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
return Path("recordings") / f"{started_at}_events.csv"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Parse the ESP32 encoder binary stream into an event CSV."
|
||||
)
|
||||
source = parser.add_mutually_exclusive_group()
|
||||
source.add_argument("--port", help="Serial port to read from.")
|
||||
source.add_argument("--input", type=Path, help="Captured binary stream to parse.")
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output",
|
||||
type=Path,
|
||||
default=default_output_path(),
|
||||
help="Output CSV path.",
|
||||
)
|
||||
parser.add_argument("--baud", type=int, default=BAUD_RATE)
|
||||
parser.add_argument(
|
||||
"--duration",
|
||||
type=float,
|
||||
help="Recording duration in seconds. Without this, serial reads until Ctrl-C.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
events = EventTable()
|
||||
|
||||
try:
|
||||
if args.input is not None:
|
||||
with args.input.open("rb") as stream:
|
||||
frames = parse_stream(stream, events)
|
||||
else:
|
||||
import serial
|
||||
|
||||
port = args.port or choose_port()
|
||||
if port is None:
|
||||
return 0
|
||||
|
||||
try:
|
||||
with serial.Serial(port, args.baud, timeout=0.1) as stream:
|
||||
stream.reset_input_buffer()
|
||||
print(f"Reading {port} at {args.baud} baud. Press Ctrl-C to stop.")
|
||||
frames = parse_stream(
|
||||
stream,
|
||||
events,
|
||||
args.duration,
|
||||
eof_on_empty=False,
|
||||
)
|
||||
except serial.SerialException as exc:
|
||||
print(f"Serial error: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
except KeyboardInterrupt:
|
||||
print()
|
||||
frames = len(events.rows)
|
||||
except ImportError as exc:
|
||||
print(f"Missing Python dependency: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
events.write_csv(args.output)
|
||||
print(f"Wrote {len(events.rows)} event times from {frames} frames to {args.output}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user