284 lines
8.1 KiB
Python
284 lines
8.1 KiB
Python
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())
|