From 7af5477f044950dbee36e35531346a96a92896ba Mon Sep 17 00:00:00 2001 From: Mathew Beane Date: Fri, 2 Oct 2026 22:01:22 -0400 Subject: [PATCH] fix(scripts): parse the ADR-018 header layout in the UDP capture scripts Both capture scripts decoded a header layout the firmware does not send. csi_collector.c and esp32_parser.rs use: magic u32 @0, node_id @4, n_antennas @5, n_subcarriers u16 @6, freq_mhz u32 @8, sequence u32 @12, rssi i8 @16, noise_floor i8 @17, ppdu_type @18, flags @19, then n_antennas * n_subcarriers int8 I/Q pairs from @20. - collect-training-data.py read a 16-byte header and int16 I/Q, so its length check rejected every real single-antenna frame. It now reads the 20-byte header and int8 pairs, records freq/channel/seq/ppdu/flags, and lets the recorder stamp receive time (the header has no timestamp). - record-csi-udp.py read rssi from byte 6 and channel from byte 7 (the n_subcarriers bytes) and treated bytes 8-19 as I/Q, so a 64-subcarrier frame came out as 70 subcarriers with header bytes in the first six amplitudes. It now checks the magic, reads the real fields, and writes only the I/Q block to iq_hex. - Add scripts/tests/test_capture_adr018.py, which builds frames with the esp32_parser.rs test layout and checks the decoded fields. Co-Authored-By: claude-flow --- scripts/collect-training-data.py | 51 +++++-- scripts/record-csi-udp.py | 54 +++++-- scripts/tests/test_capture_adr018.py | 220 +++++++++++++++++++++++++++ 3 files changed, 297 insertions(+), 28 deletions(-) create mode 100644 scripts/tests/test_capture_adr018.py diff --git a/scripts/collect-training-data.py b/scripts/collect-training-data.py index 6b2e4d6b0..475bddefb 100644 --- a/scripts/collect-training-data.py +++ b/scripts/collect-training-data.py @@ -61,9 +61,25 @@ MAGIC_FEATURES = 0xC5110003 FEATURE_PKT_FMT = " int: + """802.11 channel number for a 2.4/5 GHz centre frequency (0 if unknown).""" + if freq_mhz == 2484: + return 14 + if 2412 <= freq_mhz <= 2472: + return (freq_mhz - 2407) // 5 + if 5000 <= freq_mhz <= 5900: + return (freq_mhz - 5000) // 5 + return 0 # ── Packet parsing ─────────────────────────────────────────────────────────── @@ -114,22 +130,22 @@ def _parse_feature_packet(data: bytes) -> Optional[dict]: def _parse_raw_csi_packet(data: bytes) -> Optional[dict]: """Parse ADR-018 raw CSI frame with full subcarrier data.""" try: - magic, node_id, ant_cfg, n_sub, rssi, noise, channel, ts_ms = struct.unpack_from( - RAW_CSI_HDR_FMT, data - ) + (magic, node_id, n_ant, n_sub, freq_mhz, seq, + rssi, noise, ppdu_type, flags) = struct.unpack_from(RAW_CSI_HDR_FMT, data) except struct.error: return None if magic != MAGIC_CSI_RAW: return None - - # Subcarrier data follows header as int16 I/Q pairs - payload_offset = RAW_CSI_HDR_SIZE - expected_bytes = n_sub * 2 * 2 # n_sub * (I + Q) * int16 - if len(data) < payload_offset + expected_bytes: + if not 1 <= n_ant <= RAW_CSI_MAX_ANTENNAS or n_sub > RAW_CSI_MAX_SUBCARRIERS: return None - iq_data = struct.unpack_from(f"<{n_sub * 2}h", data, payload_offset) + # I/Q data follows the header as int8 pairs, antenna-major + n_pairs = n_ant * n_sub + if len(data) < RAW_CSI_HDR_SIZE + n_pairs * 2: + return None + + iq_data = struct.unpack_from(f"<{n_pairs * 2}b", data, RAW_CSI_HDR_SIZE) # Convert I/Q pairs to amplitude subcarriers = [] for i in range(0, len(iq_data), 2): @@ -137,13 +153,18 @@ def _parse_raw_csi_packet(data: bytes) -> Optional[dict]: amplitude = (real ** 2 + imag ** 2) ** 0.5 subcarriers.append(amplitude) + # The header has no timestamp; CsiRecorder stamps the receive time. return { "type": "raw_csi", "node_id": node_id, - "antenna_config": ant_cfg, + "antenna_config": n_ant, + "n_antennas": n_ant, "n_subcarriers": n_sub, - "channel": channel, - "timestamp": ts_ms / 1000.0, + "freq_mhz": freq_mhz, + "channel": freq_to_channel(freq_mhz), + "seq": seq, + "ppdu_type": ppdu_type, + "flags": flags, "subcarriers": subcarriers, "rssi": float(rssi), "noise_floor": float(noise), diff --git a/scripts/record-csi-udp.py b/scripts/record-csi-udp.py index 1f74b27bf..4d48f5cb5 100644 --- a/scripts/record-csi-udp.py +++ b/scripts/record-csi-udp.py @@ -18,24 +18,45 @@ import time from datetime import datetime, timezone +# ADR-018 header, as written by csi_collector.c and read by esp32_parser.rs: +# magic u32 @0, node_id @4, n_antennas @5, n_subcarriers u16 @6, freq_mhz u32 @8, +# sequence u32 @12, rssi i8 @16, noise_floor i8 @17, ppdu_type @18, flags @19, +# then n_antennas * n_subcarriers int8 I/Q pairs from @20. +CSI_MAGIC = 0xC5110001 +CSI_HDR_FMT = " 4 else 0 - rssi = struct.unpack('b', bytes([data[6]]))[0] if len(data) > 6 else 0 - channel = data[7] if len(data) > 7 else 0 + (magic, node_id, n_antennas, n_sub, freq_mhz, seq, + rssi, noise_floor, ppdu_type, flags) = struct.unpack_from(CSI_HDR_FMT, data) + # Same antenna/subcarrier limits as esp32_parser.rs + if magic != CSI_MAGIC or not 1 <= n_antennas <= 4 or n_sub > 256: + return None - # IQ data starts at offset 8 - iq_data = data[8:] if len(data) > 8 else b'' - n_subcarriers = len(iq_data) // 2 # I,Q pairs + n_pairs = n_antennas * n_sub + iq_data = data[CSI_HDR_SIZE:CSI_HDR_SIZE + n_pairs * 2] + if len(iq_data) < n_pairs * 2: + return None # Compute amplitudes amplitudes = [] - for i in range(0, len(iq_data) - 1, 2): + for i in range(0, len(iq_data), 2): I = struct.unpack('b', bytes([iq_data[i]]))[0] Q = struct.unpack('b', bytes([iq_data[i + 1]]))[0] amplitudes.append(round((I * I + Q * Q) ** 0.5, 2)) @@ -46,9 +67,16 @@ def parse_csi_packet(data): "timestamp": datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z"), "ts_ns": time.time_ns(), "node_id": node_id, + "seq": seq, "rssi": rssi, - "channel": channel, - "subcarriers": n_subcarriers, + "noise_floor": noise_floor, + "freq_mhz": freq_mhz, + "channel": freq_to_channel(freq_mhz), + "ppdu_type": ppdu_type, + "flags": flags, + "n_antennas": n_antennas, + "n_subcarriers": n_sub, + "subcarriers": n_pairs, # I/Q pairs in iq_hex (n_antennas * n_subcarriers) "amplitudes": amplitudes, "iq_hex": iq_data.hex(), } diff --git a/scripts/tests/test_capture_adr018.py b/scripts/tests/test_capture_adr018.py new file mode 100644 index 000000000..b8715e08d --- /dev/null +++ b/scripts/tests/test_capture_adr018.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +"""ADR-018 header parsing in the UDP capture scripts. + +Builds real ADR-018 frames with the layout used by +v2/crates/wifi-densepose-hardware/src/esp32_parser.rs (build_test_frame) and +firmware/esp32-csi-node/main/csi_collector.c, then checks the fields that +scripts/collect-training-data.py and scripts/record-csi-udp.py decode. + +Run: python -m pytest scripts/tests/test_capture_adr018.py -q +""" + +from __future__ import annotations + +import importlib.util +import json +import math +import struct +import time +from pathlib import Path + +import pytest + +SCRIPTS_DIR = Path(__file__).resolve().parents[1] + +ADR018_MAGIC = 0xC5110001 +FEATURE_MAGIC = 0xC5110003 + + +def load_script(filename: str): + """Import a hyphenated script from scripts/ as a module.""" + name = filename.replace("-", "_").removesuffix(".py") + spec = importlib.util.spec_from_file_location(name, SCRIPTS_DIR / filename) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +collect = load_script("collect-training-data.py") +record = load_script("record-csi-udp.py") + + +def build_frame( + node_id: int, + n_antennas: int, + pairs: list[tuple[int, int]], + *, + freq_mhz: int = 2437, + seq: int = 1, + rssi: int = -50, + noise: int = -95, + ppdu: int = 0, + flags: int = 0, + magic: int = ADR018_MAGIC, +) -> bytes: + """ADR-018 frame, same byte layout as esp32_parser.rs build_test_frame_with_he. + + 0 magic u32 | 4 node_id | 5 n_antennas | 6 n_subcarriers u16 | 8 freq_mhz u32 + | 12 sequence u32 | 16 rssi i8 | 17 noise i8 | 18 ppdu | 19 flags | 20 I/Q i8 pairs + """ + n_sub = len(pairs) // n_antennas if n_antennas else len(pairs) + buf = bytearray() + buf += struct.pack(" amplitude 5, last pair (-6, 8) -> 10. +PAIRS_64 = [(3, 4)] + [(1, -1)] * 62 + [(-6, 8)] +# Sequence chosen so bytes 14/15 (seq high bytes) differ from rssi/noise. +SEQ = 0x01020304 + + +def test_builder_matches_rust_header_size(): + frame = build_frame(1, 1, [(0, 0)] * 64) + assert len(frame) == 20 + 64 * 2 # 148 bytes, the HT20 frame size in #1005 + + +# ── collect-training-data.py ──────────────────────────────────────────────── + +class TestCollectTrainingData: + def test_single_antenna_header_fields(self): + frame = build_frame(7, 1, PAIRS_64, seq=SEQ, rssi=-50, noise=-95, + ppdu=1, flags=0x01) + out = collect.parse_packet(frame) + assert out is not None + assert out["type"] == "raw_csi" + assert out["node_id"] == 7 + assert out["antenna_config"] == 1 + assert out["n_antennas"] == 1 + assert out["n_subcarriers"] == 64 + assert out["freq_mhz"] == 2437 + assert out["channel"] == 6 + assert out["seq"] == SEQ + assert out["rssi"] == -50.0 + assert out["noise_floor"] == -95.0 + assert out["ppdu_type"] == 1 + assert out["flags"] == 0x01 + + def test_single_antenna_amplitudes_from_int8_pairs(self): + out = collect.parse_packet(build_frame(7, 1, PAIRS_64)) + amps = out["subcarriers"] + assert len(amps) == 64 + assert amps[0] == pytest.approx(5.0) + assert amps[1] == pytest.approx(math.sqrt(2)) + assert amps[-1] == pytest.approx(10.0) + + def test_multi_antenna_reads_all_pairs(self): + pairs = [(3, 4)] * 4 + [(6, 8)] * 4 # antenna 0, then antenna 1 + out = collect.parse_packet(build_frame(2, 2, pairs)) + assert out["n_antennas"] == 2 + assert out["n_subcarriers"] == 4 + assert out["subcarriers"] == pytest.approx([5.0] * 4 + [10.0] * 4) + + def test_5ghz_channel(self): + out = collect.parse_packet(build_frame(1, 1, PAIRS_64, freq_mhz=5180)) + assert out["freq_mhz"] == 5180 + assert out["channel"] == 36 + + def test_no_header_timestamp_so_recorder_uses_receive_time(self, tmp_path): + out = collect.parse_packet(build_frame(3, 1, PAIRS_64, seq=SEQ)) + assert "timestamp" not in out + rec = collect.CsiRecorder(str(tmp_path), "t", "walking") + rec.open() + before = time.time() + rec.write_frame(out) + rec.close() + line = json.loads(rec.file_path.read_text().splitlines()[0]) + assert before - 1 <= line["timestamp"] <= time.time() + 1 + assert line["rssi"] == -50.0 + assert line["features"]["seq"] == SEQ + assert len(line["subcarriers"]) == 64 + + def test_trailing_bytes_ignored(self): + out = collect.parse_packet(build_frame(1, 1, PAIRS_64) + b"\xff" * 8) + assert len(out["subcarriers"]) == 64 + + def test_truncated_iq_rejected(self): + assert collect.parse_packet(build_frame(1, 1, PAIRS_64)[:-1]) is None + + def test_short_header_rejected(self): + assert collect.parse_packet(build_frame(1, 1, PAIRS_64)[:19]) is None + + def test_zero_antennas_rejected(self): + frame = bytearray(build_frame(1, 1, PAIRS_64)) + frame[5] = 0 + assert collect.parse_packet(bytes(frame)) is None + + def test_feature_packet_still_parses(self): + feats = [0.1 * i for i in range(8)] + pkt = struct.pack("