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 <ruv@ruv.net>

scripts/collect-training-data.py

@@ -61,9 +61,25 @@ MAGIC_FEATURES = 0xC5110003
 FEATURE_PKT_FMT = "<IBBHq8f"
 FEATURE_PKT_SIZE = struct.calcsize(FEATURE_PKT_FMT)  # 48 bytes
 
-# Raw CSI header: magic(4) + node_id(1) + antenna_cfg(1) + n_sub(2) + rssi(1) + noise(1) + channel(1) + reserved(1) + timestamp_ms(4)
-RAW_CSI_HDR_FMT = "<IBBHbbBxI"
-RAW_CSI_HDR_SIZE = struct.calcsize(RAW_CSI_HDR_FMT)  # 16 bytes
+# Raw CSI header, as written by csi_collector.c and read by esp32_parser.rs:
+# magic(4) + node_id(1) + n_antennas(1) + n_subcarriers(2) + freq_mhz(4) +
+# sequence(4) + rssi(1) + noise_floor(1) + ppdu_type(1) + flags(1), then
+# n_antennas * n_subcarriers int8 I/Q pairs.
+RAW_CSI_HDR_FMT = "<IBBHIIbbBB"
+RAW_CSI_HDR_SIZE = struct.calcsize(RAW_CSI_HDR_FMT)  # 20 bytes
+RAW_CSI_MAX_ANTENNAS = 4
+RAW_CSI_MAX_SUBCARRIERS = 256
+
+
+def freq_to_channel(freq_mhz: int) -> 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),

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 = "<IBBHIIbbBB"
+CSI_HDR_SIZE = struct.calcsize(CSI_HDR_FMT)  # 20 bytes
+
+
+def freq_to_channel(freq_mhz):
+    """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
+
+
 def parse_csi_packet(data):
-    """Parse ADR-018 binary CSI packet into dict."""
-    if len(data) < 8:
+    """Parse ADR-018 binary CSI packet into dict, or None if it is not one."""
+    if len(data) < CSI_HDR_SIZE:
         return None
 
-    # ADR-018 header: [magic(2), len(2), node_id(1), seq(1), rssi(1), channel(1), iq_data...]
-    # Simplified: extract what we can from the raw packet
-    node_id = data[4] if len(data) > 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(),
     }

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("<I", magic)
+    buf.append(node_id)
+    buf.append(n_antennas)
+    buf += struct.pack("<H", n_sub)
+    buf += struct.pack("<I", freq_mhz)
+    buf += struct.pack("<I", seq)
+    buf += struct.pack("<b", rssi)
+    buf += struct.pack("<b", noise)
+    buf.append(ppdu)
+    buf.append(flags)
+    for i, q in pairs:
+        buf += struct.pack("<bb", i, q)
+    return bytes(buf)
+
+
+# 64 subcarriers, first pair (3, 4) -> 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("<IBBHq8f", FEATURE_MAGIC, 4, 0, 9, 123456, *feats)
+        out = collect.parse_packet(pkt)
+        assert out["type"] == "features"
+        assert out["node_id"] == 4
+        assert out["seq"] == 9
+        assert out["features"] == pytest.approx(feats)
+
+
+# ── record-csi-udp.py ───────────────────────────────────────────────────────
+
+class TestRecordCsiUdp:
+    def test_single_antenna_header_fields(self):
+        frame = build_frame(7, 1, PAIRS_64, seq=SEQ, rssi=-50, noise=-95)
+        out = record.parse_csi_packet(frame)
+        assert out is not None
+        assert out["node_id"] == 7
+        assert out["rssi"] == -50
+        assert out["noise_floor"] == -95
+        assert out["freq_mhz"] == 2437
+        assert out["channel"] == 6
+        assert out["seq"] == SEQ
+        assert out["n_antennas"] == 1
+        assert out["n_subcarriers"] == 64
+
+    def test_iq_payload_starts_at_offset_20(self):
+        frame = build_frame(7, 1, PAIRS_64)
+        out = record.parse_csi_packet(frame)
+        assert out["subcarriers"] == 64
+        assert out["iq_hex"] == frame[20:].hex()
+        assert len(out["amplitudes"]) == 64
+        assert out["amplitudes"][0] == 5.0
+        assert out["amplitudes"][-1] == 10.0
+
+    def test_multi_antenna(self):
+        pairs = [(3, 4)] * 4 + [(6, 8)] * 4
+        out = record.parse_csi_packet(build_frame(2, 2, pairs))
+        assert out["n_antennas"] == 2
+        assert out["n_subcarriers"] == 4
+        assert out["subcarriers"] == 8
+        assert out["amplitudes"] == [5.0] * 4 + [10.0] * 4
+
+    def test_trailing_bytes_not_treated_as_iq(self):
+        frame = build_frame(1, 1, PAIRS_64)
+        out = record.parse_csi_packet(frame + b"\x7f" * 6)
+        assert out["iq_hex"] == frame[20:].hex()
+
+    def test_other_magic_rejected(self):
+        feats = [0.0] * 8
+        pkt = struct.pack("<IBBHq8f", FEATURE_MAGIC, 4, 0, 9, 123456, *feats)
+        assert record.parse_csi_packet(pkt) is None
+
+    def test_truncated_iq_rejected(self):
+        assert record.parse_csi_packet(build_frame(1, 1, PAIRS_64)[:-1]) is None
+
+    def test_short_header_rejected(self):
+        assert record.parse_csi_packet(build_frame(1, 1, PAIRS_64)[:19]) is None
+
+    def test_timestamps_present(self):
+        out = record.parse_csi_packet(build_frame(1, 1, PAIRS_64))
+        assert out["timestamp"].endswith("Z")
+        assert isinstance(out["ts_ns"], int)