starred/RuView
1
0 Fork 0
mirror of https://github.com/ruvnet/RuView.git synced 2026-10-18 21:00:22 -05:00
Code Issues Packages Projects Releases 10 Wiki Activity
Files
7af5477f044950dbee36e35531346a96a92896ba
RuView/scripts/collect-training-data.py
T
Mathew Beaneandclaude-flow 7af5477f04 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>
2026-10-02 22:01:22 -04:00

505 lines
17 KiB
Python
Raw Blame History

#!/usr/bin/env python3
"""
WiFi-DensePose Training Data Collector
Listens on UDP for CSI data from ESP32 nodes and records to .csi.jsonl
files compatible with the Rust training pipeline (MmFiDataset / CsiDataset).
Supports two packet formats:
- ADR-069 feature vectors (magic 0xC5110003, 48 bytes) — 8-dim pre-extracted
- ADR-018 raw CSI frames (magic 0xC5110001, variable) — full subcarrier data
Usage:
# Interactive — prompts for scenario labels
python scripts/collect-training-data.py --port 5006
# Scripted — fixed label, 60s per recording
python scripts/collect-training-data.py --port 5006 --label walking --duration 60
# Multiple scenarios in sequence
python scripts/collect-training-data.py --port 5006 --scenarios walking,standing,sitting --duration 30
# Dual-node collection (two ESP32s on different ports)
python scripts/collect-training-data.py --port 5005 --port2 5006 --label walking
# Generate manifest only from existing recordings
python scripts/collect-training-data.py --manifest-only --output-dir data/recordings
Prerequisites:
- ESP32 nodes streaming CSI on UDP (see firmware/esp32-csi-node)
- Python 3.9+
"""
from __future__ import annotations
import argparse
import json
import logging
import os
import socket
import struct
import sys
import time
import signal
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
datefmt="%H:%M:%S",
)
log = logging.getLogger("collect-data")
# ── Packet formats (must match firmware) ─────────────────────────────────────
# ADR-018 raw CSI frame header
MAGIC_CSI_RAW = 0xC5110001
# ADR-069 feature vector packet
MAGIC_FEATURES = 0xC5110003
FEATURE_PKT_FMT = "<IBBHq8f"
FEATURE_PKT_SIZE = struct.calcsize(FEATURE_PKT_FMT) # 48 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 ───────────────────────────────────────────────────────────
def parse_packet(data: bytes) -> Optional[dict]:
"""Parse a UDP packet into a frame dict, or None if unrecognized."""
if len(data) < 4:
return None
magic = struct.unpack_from("<I", data)[0]
if magic == MAGIC_FEATURES and len(data) >= FEATURE_PKT_SIZE:
return _parse_feature_packet(data)
elif magic == MAGIC_CSI_RAW and len(data) >= RAW_CSI_HDR_SIZE:
return _parse_raw_csi_packet(data)
else:
return None
def _parse_feature_packet(data: bytes) -> Optional[dict]:
"""Parse ADR-069 feature vector packet (48 bytes)."""
try:
magic, node_id, _, seq, ts_us, *features = struct.unpack_from(FEATURE_PKT_FMT, data)
except struct.error:
return None
if magic != MAGIC_FEATURES:
return None
# Reject NaN/inf
import math
if any(math.isnan(f) or math.isinf(f) for f in features):
return None
return {
"type": "features",
"node_id": node_id,
"seq": seq,
"timestamp_us": ts_us,
"timestamp": ts_us / 1_000_000.0,
"features": features,
"subcarriers": features, # Use features as subcarrier proxy for training
"rssi": 0.0,
"noise_floor": 0.0,
}
def _parse_raw_csi_packet(data: bytes) -> Optional[dict]:
"""Parse ADR-018 raw CSI frame with full subcarrier data."""
try:
(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
if not 1 <= n_ant <= RAW_CSI_MAX_ANTENNAS or n_sub > RAW_CSI_MAX_SUBCARRIERS:
return None
# 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):
real, imag = iq_data[i], iq_data[i + 1]
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": n_ant,
"n_antennas": n_ant,
"n_subcarriers": n_sub,
"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),
}
# ── JSONL recording ──────────────────────────────────────────────────────────
class CsiRecorder:
"""Records CSI frames to .csi.jsonl files compatible with the Rust pipeline."""
def __init__(self, output_dir: str, session_name: str, label: Optional[str] = None):
self.output_dir = Path(output_dir)
self.output_dir.mkdir(parents=True, exist_ok=True)
ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
safe_name = session_name.replace(" ", "_").replace("/", "_")
self.session_id = f"{safe_name}-{ts}"
self.label = label
self.file_path = self.output_dir / f"{self.session_id}.csi.jsonl"
self.meta_path = self.output_dir / f"{self.session_id}.csi.meta.json"
self.frame_count = 0
self.start_time = time.time()
self.started_at = datetime.now(timezone.utc).isoformat()
self._file = None
def open(self):
self._file = open(self.file_path, "a", encoding="utf-8")
log.info(f"Recording to: {self.file_path}")
def write_frame(self, frame: dict):
"""Write a single frame as a JSONL line."""
if self._file is None:
return
record = {
"timestamp": frame.get("timestamp", time.time()),
"subcarriers": frame.get("subcarriers", []),
"rssi": frame.get("rssi", 0.0),
"noise_floor": frame.get("noise_floor", 0.0),
"features": {
k: v for k, v in frame.items()
if k not in ("timestamp", "subcarriers", "rssi", "noise_floor", "type")
},
}
line = json.dumps(record, separators=(",", ":"))
self._file.write(line + "\n")
self.frame_count += 1
if self.frame_count % 500 == 0:
self._file.flush()
def close(self) -> dict:
"""Close the recording and write metadata. Returns session info."""
if self._file:
self._file.flush()
self._file.close()
self._file = None
ended_at = datetime.now(timezone.utc).isoformat()
elapsed = time.time() - self.start_time
file_size = self.file_path.stat().st_size if self.file_path.exists() else 0
meta = {
"id": self.session_id,
"name": self.session_id,
"label": self.label,
"started_at": self.started_at,
"ended_at": ended_at,
"duration_secs": round(elapsed, 2),
"frame_count": self.frame_count,
"file_size_bytes": file_size,
"file_path": str(self.file_path),
"fps": round(self.frame_count / elapsed, 1) if elapsed > 0 else 0,
}
with open(self.meta_path, "w", encoding="utf-8") as f:
json.dump(meta, f, indent=2)
log.info(
f"Recording stopped: {self.frame_count} frames in {elapsed:.1f}s "
f"({meta['fps']} fps, {file_size / 1024:.1f} KB)"
)
return meta
# ── Manifest generation ──────────────────────────────────────────────────────
def generate_manifest(output_dir: str) -> dict:
"""Scan recordings directory and generate a dataset manifest JSON."""
rec_dir = Path(output_dir)
sessions = []
for meta_file in sorted(rec_dir.glob("*.csi.meta.json")):
try:
with open(meta_file, "r") as f:
meta = json.load(f)
sessions.append(meta)
except (json.JSONDecodeError, OSError) as e:
log.warning(f"Skipping {meta_file}: {e}")
# Aggregate stats
total_frames = sum(s.get("frame_count", 0) for s in sessions)
total_bytes = sum(s.get("file_size_bytes", 0) for s in sessions)
labels = sorted(set(s.get("label", "unlabeled") or "unlabeled" for s in sessions))
manifest = {
"dataset": "wifi-densepose-csi",
"generated_at": datetime.now(timezone.utc).isoformat(),
"directory": str(rec_dir),
"num_sessions": len(sessions),
"total_frames": total_frames,
"total_size_bytes": total_bytes,
"total_size_mb": round(total_bytes / (1024 * 1024), 2),
"labels": labels,
"sessions": sessions,
}
manifest_path = rec_dir / "manifest.json"
with open(manifest_path, "w", encoding="utf-8") as f:
json.dump(manifest, f, indent=2)
log.info(
f"Manifest: {len(sessions)} sessions, {total_frames} frames, "
f"{manifest['total_size_mb']} MB, labels={labels}"
)
log.info(f"Written to: {manifest_path}")
return manifest
# ── UDP listener ─────────────────────────────────────────────────────────────
def collect_session(
port: int,
port2: Optional[int],
output_dir: str,
label: str,
duration: float,
session_name: Optional[str] = None,
) -> dict:
"""Run a single collection session. Returns session metadata."""
name = session_name or label or "session"
recorder = CsiRecorder(output_dir, name, label)
recorder.open()
# Bind primary socket
sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(("0.0.0.0", port))
sock.settimeout(1.0)
sockets = [sock]
# Bind secondary socket if specified
if port2:
sock2 = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock2.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock2.bind(("0.0.0.0", port2))
sock2.settimeout(0.1)
sockets.append(sock2)
log.info(
f"Collecting '{label}' for {duration}s on port(s) "
f"{port}{f', {port2}' if port2 else ''}"
)
start = time.time()
dropped = 0
try:
while time.time() - start < duration:
for s in sockets:
try:
data, addr = s.recvfrom(4096)
except socket.timeout:
continue
frame = parse_packet(data)
if frame:
recorder.write_frame(frame)
else:
dropped += 1
# Progress update every 5s
elapsed = time.time() - start
if recorder.frame_count > 0 and int(elapsed) % 5 == 0 and int(elapsed) > 0:
remaining = duration - elapsed
if remaining > 0 and int(elapsed * 10) % 50 == 0:
log.info(
f" {recorder.frame_count} frames collected, "
f"{remaining:.0f}s remaining..."
)
except KeyboardInterrupt:
log.info("Interrupted by user.")
finally:
for s in sockets:
s.close()
if dropped > 0:
log.warning(f" {dropped} unrecognized packets dropped")
return recorder.close()
# ── Main ─────────────────────────────────────────────────────────────────────
def main():
parser = argparse.ArgumentParser(
description="Collect CSI training data from ESP32 nodes via UDP",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Interactive label input
python scripts/collect-training-data.py --port 5006
# Fixed label, 60 seconds
python scripts/collect-training-data.py --port 5006 --label walking --duration 60
# Multiple scenarios
python scripts/collect-training-data.py --port 5006 --scenarios walking,standing,sitting --duration 30
# Dual ESP32 nodes
python scripts/collect-training-data.py --port 5005 --port2 5006 --label test
# Generate manifest from existing recordings
python scripts/collect-training-data.py --manifest-only
""",
)
parser.add_argument("--port", type=int, default=5006, help="Primary UDP port (default: 5006)")
parser.add_argument("--port2", type=int, default=None, help="Secondary UDP port for dual-node")
parser.add_argument("--output-dir", default="data/recordings", help="Output directory (default: data/recordings)")
parser.add_argument("--label", default=None, help="Activity label for the recording")
parser.add_argument("--duration", type=float, default=30.0, help="Recording duration in seconds (default: 30)")
parser.add_argument("--scenarios", default=None, help="Comma-separated list of scenarios to record sequentially")
parser.add_argument("--pause", type=float, default=5.0, help="Pause between scenarios in seconds (default: 5)")
parser.add_argument("--manifest-only", action="store_true", help="Only generate manifest from existing recordings")
parser.add_argument("--repeats", type=int, default=1, help="Number of repeats per scenario (default: 1)")
args = parser.parse_args()
# Manifest-only mode
if args.manifest_only:
generate_manifest(args.output_dir)
return
# Collect scenarios
all_sessions = []
if args.scenarios:
# Multi-scenario sequential collection
scenarios = [s.strip() for s in args.scenarios.split(",") if s.strip()]
total = len(scenarios) * args.repeats
idx = 0
for repeat in range(args.repeats):
for scenario in scenarios:
idx += 1
print(f"\n{'='*60}")
print(f" Scenario {idx}/{total}: '{scenario}' (repeat {repeat+1}/{args.repeats})")
print(f" Duration: {args.duration}s")
print(f"{'='*60}")
if idx > 1:
print(f" Starting in {args.pause}s... (get into position)")
time.sleep(args.pause)
meta = collect_session(
port=args.port,
port2=args.port2,
output_dir=args.output_dir,
label=scenario,
duration=args.duration,
session_name=f"{scenario}_r{repeat+1:02d}",
)
all_sessions.append(meta)
elif args.label:
# Single labeled recording
meta = collect_session(
port=args.port,
port2=args.port2,
output_dir=args.output_dir,
label=args.label,
duration=args.duration,
)
all_sessions.append(meta)
else:
# Interactive mode — prompt for labels
print("\nInteractive data collection mode.")
print("Type a label for each recording, or 'q' to quit.\n")
while True:
label = input("Label (or 'q' to quit): ").strip()
if label.lower() in ("q", "quit", "exit"):
break
if not label:
print(" Empty label. Try again.")
continue
duration = args.duration
try:
dur_input = input(f"Duration in seconds [{duration}]: ").strip()
if dur_input:
duration = float(dur_input)
except ValueError:
pass
print(f" Recording '{label}' for {duration}s — starting now...")
meta = collect_session(
port=args.port,
port2=args.port2,
output_dir=args.output_dir,
label=label,
duration=duration,
)
all_sessions.append(meta)
print()
# Generate manifest
if all_sessions:
print(f"\nCollected {len(all_sessions)} session(s).")
manifest = generate_manifest(args.output_dir)
total_frames = sum(s.get("frame_count", 0) for s in all_sessions)
print(f"\nSummary:")
print(f" Sessions: {len(all_sessions)}")
print(f" Total frames: {total_frames}")
print(f" Output: {args.output_dir}/")
print(f" Manifest: {args.output_dir}/manifest.json")
else:
print("No sessions recorded.")
if __name__ == "__main__":
main()
Reference in new issue View Git Blame Copy Permalink