use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::io;
use std::path::{Path, PathBuf};
#[derive(Debug)]
pub enum DatasetError {
Io(io::Error),
Format(String),
Missing(String),
Shape(String),
}
impl fmt::Display for DatasetError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Io(e) => write!(f, "I/O error: {e}"),
Self::Format(s) => write!(f, "format error: {s}"),
Self::Missing(s) => write!(f, "missing: {s}"),
Self::Shape(s) => write!(f, "shape error: {s}"),
}
}
}
impl std::error::Error for DatasetError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
if let Self::Io(e) = self {
Some(e)
} else {
None
}
}
}
impl From<io::Error> for DatasetError {
fn from(e: io::Error) -> Self {
Self::Io(e)
}
}
pub type Result<T> = std::result::Result<T, DatasetError>;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NpyArray {
pub shape: Vec<usize>,
pub data: Vec<f32>,
}
impl NpyArray {
pub fn len(&self) -> usize {
self.data.len()
}
pub fn is_empty(&self) -> bool {
self.data.is_empty()
}
pub fn ndim(&self) -> usize {
self.shape.len()
}
}
pub struct NpyReader;
impl NpyReader {
pub fn read_file(path: &Path) -> Result<NpyArray> {
Self::parse(&std::fs::read(path)?)
}
pub fn parse(buf: &[u8]) -> Result<NpyArray> {
if buf.len() < 10 {
return Err(DatasetError::Format("file too small for .npy".into()));
}
if &buf[0..6] != b"\x93NUMPY" {
return Err(DatasetError::Format("missing .npy magic".into()));
}
let major = buf[6];
let (header_len, header_start) = match major {
1 => (u16::from_le_bytes([buf[8], buf[9]]) as usize, 10usize),
2 | 3 => {
if buf.len() < 12 {
return Err(DatasetError::Format("truncated v2 header".into()));
}
(
u32::from_le_bytes([buf[8], buf[9], buf[10], buf[11]]) as usize,
12,
)
}
_ => {
return Err(DatasetError::Format(format!(
"unsupported .npy version {major}"
)))
}
};
let header_end = header_start + header_len;
if header_end > buf.len() {
return Err(DatasetError::Format("header past EOF".into()));
}
let hdr = std::str::from_utf8(&buf[header_start..header_end])
.map_err(|_| DatasetError::Format("non-UTF8 header".into()))?;
let dtype = Self::extract_field(hdr, "descr")?;
let is_f64 = dtype.contains("f8") || dtype.contains("float64");
let is_f32 = dtype.contains("f4") || dtype.contains("float32");
let is_big = dtype.starts_with('>');
if !is_f32 && !is_f64 {
return Err(DatasetError::Format(format!("unsupported dtype '{dtype}'")));
}
let fortran = Self::extract_field(hdr, "fortran_order")
.unwrap_or_else(|_| "False".into())
.contains("True");
let shape = Self::parse_shape(hdr)?;
let elem_sz: usize = if is_f64 { 8 } else { 4 };
let total: usize = shape.iter().product::<usize>().max(1);
if header_end + total * elem_sz > buf.len() {
return Err(DatasetError::Format("data truncated".into()));
}
let raw = &buf[header_end..header_end + total * elem_sz];
let mut data: Vec<f32> = if is_f64 {
raw.chunks_exact(8)
.map(|c| {
let v = if is_big {
f64::from_be_bytes(c.try_into().unwrap())
} else {
f64::from_le_bytes(c.try_into().unwrap())
};
v as f32
})
.collect()
} else {
raw.chunks_exact(4)
.map(|c| {
if is_big {
f32::from_be_bytes(c.try_into().unwrap())
} else {
f32::from_le_bytes(c.try_into().unwrap())
}
})
.collect()
};
if fortran && shape.len() == 2 {
let (r, c) = (shape[0], shape[1]);
let mut cd = vec![0.0f32; data.len()];
for ri in 0..r {
for ci in 0..c {
cd[ri * c + ci] = data[ci * r + ri];
}
}
data = cd;
}
let shape = if shape.is_empty() { vec![1] } else { shape };
Ok(NpyArray { shape, data })
}
fn extract_field(hdr: &str, field: &str) -> Result<String> {
for pat in &[
format!("'{field}': "),
format!("'{field}':"),
format!("\"{field}\": "),
] {
if let Some(s) = hdr.find(pat.as_str()) {
let rest = &hdr[s + pat.len()..];
let end = rest
.find(',')
.or_else(|| rest.find('}'))
.unwrap_or(rest.len());
return Ok(rest[..end]
.trim()
.trim_matches('\'')
.trim_matches('"')
.into());
}
}
Err(DatasetError::Format(format!("field '{field}' not found")))
}
fn parse_shape(hdr: &str) -> Result<Vec<usize>> {
let si = hdr
.find("'shape'")
.or_else(|| hdr.find("\"shape\""))
.ok_or_else(|| DatasetError::Format("no 'shape'".into()))?;
let rest = &hdr[si..];
let ps = rest
.find('(')
.ok_or_else(|| DatasetError::Format("no '('".into()))?;
let pe = rest[ps..]
.find(')')
.ok_or_else(|| DatasetError::Format("no ')'".into()))?;
let inner = rest[ps + 1..ps + pe].trim();
if inner.is_empty() {
return Ok(vec![]);
}
inner
.split(',')
.map(|s| s.trim())
.filter(|s| !s.is_empty())
.map(|s| {
s.parse::<usize>()
.map_err(|_| DatasetError::Format(format!("bad dim: '{s}'")))
})
.collect()
}
}
pub struct MatReader;
const MI_INT8: u32 = 1;
#[allow(dead_code)]
const MI_UINT8: u32 = 2;
#[allow(dead_code)]
const MI_INT16: u32 = 3;
#[allow(dead_code)]
const MI_UINT16: u32 = 4;
const MI_INT32: u32 = 5;
const MI_UINT32: u32 = 6;
const MI_SINGLE: u32 = 7;
const MI_DOUBLE: u32 = 9;
const MI_MATRIX: u32 = 14;
impl MatReader {
pub fn read_file(path: &Path) -> Result<HashMap<String, NpyArray>> {
Self::parse(&std::fs::read(path)?)
}
pub fn parse(buf: &[u8]) -> Result<HashMap<String, NpyArray>> {
if buf.len() < 128 {
return Err(DatasetError::Format("too small for .mat v5".into()));
}
let swap = u16::from_le_bytes([buf[126], buf[127]]) == 0x4D49;
let mut result = HashMap::new();
let mut off = 128;
while off + 8 <= buf.len() {
let (dt, ds, ts) = Self::read_tag(buf, off, swap)?;
let el_start = off + ts;
let el_end = el_start + ds;
if el_end > buf.len() {
break;
}
if dt == MI_MATRIX {
if let Ok((n, a)) = Self::parse_matrix(&buf[el_start..el_end], swap) {
result.insert(n, a);
}
}
off = (el_end + 7) & !7;
}
Ok(result)
}
fn read_tag(buf: &[u8], off: usize, swap: bool) -> Result<(u32, usize, usize)> {
if off + 4 > buf.len() {
return Err(DatasetError::Format("truncated tag".into()));
}
let raw = Self::u32(buf, off, swap);
let upper = (raw >> 16) & 0xFFFF;
if upper != 0 && upper <= 4 {
return Ok((raw & 0xFFFF, upper as usize, 4));
}
if off + 8 > buf.len() {
return Err(DatasetError::Format("truncated tag".into()));
}
Ok((raw, Self::u32(buf, off + 4, swap) as usize, 8))
}
fn parse_matrix(buf: &[u8], swap: bool) -> Result<(String, NpyArray)> {
let (mut name, mut shape, mut data) = (String::new(), Vec::new(), Vec::new());
let mut off = 0;
while off + 4 <= buf.len() {
let (st, ss, ts) = Self::read_tag(buf, off, swap)?;
let ss_start = off + ts;
let ss_end = (ss_start + ss).min(buf.len());
match st {
MI_UINT32 if shape.is_empty() && ss == 8 => {}
MI_INT32 if shape.is_empty() => {
for i in 0..ss / 4 {
shape.push(Self::i32(buf, ss_start + i * 4, swap) as usize);
}
}
MI_INT8 if name.is_empty() && ss_end <= buf.len() => {
name = String::from_utf8_lossy(&buf[ss_start..ss_end])
.trim_end_matches('\0')
.to_string();
}
MI_DOUBLE => {
for i in 0..ss / 8 {
let p = ss_start + i * 8;
if p + 8 <= buf.len() {
data.push(Self::f64(buf, p, swap) as f32);
}
}
}
MI_SINGLE => {
for i in 0..ss / 4 {
let p = ss_start + i * 4;
if p + 4 <= buf.len() {
data.push(Self::f32(buf, p, swap));
}
}
}
_ => {}
}
off = (ss_end + 7) & !7;
}
if name.is_empty() {
name = "unnamed".into();
}
if shape.is_empty() && !data.is_empty() {
shape = vec![data.len()];
}
if shape.len() == 2 {
let (r, c) = (shape[0], shape[1]);
if r * c == data.len() {
let mut cd = vec![0.0f32; data.len()];
for ri in 0..r {
for ci in 0..c {
cd[ri * c + ci] = data[ci * r + ri];
}
}
data = cd;
}
}
Ok((name, NpyArray { shape, data }))
}
fn u32(b: &[u8], o: usize, s: bool) -> u32 {
let v = [b[o], b[o + 1], b[o + 2], b[o + 3]];
if s {
u32::from_be_bytes(v)
} else {
u32::from_le_bytes(v)
}
}
fn i32(b: &[u8], o: usize, s: bool) -> i32 {
let v = [b[o], b[o + 1], b[o + 2], b[o + 3]];
if s {
i32::from_be_bytes(v)
} else {
i32::from_le_bytes(v)
}
}
fn f64(b: &[u8], o: usize, s: bool) -> f64 {
let v: [u8; 8] = b[o..o + 8].try_into().unwrap();
if s {
f64::from_be_bytes(v)
} else {
f64::from_le_bytes(v)
}
}
fn f32(b: &[u8], o: usize, s: bool) -> f32 {
let v = [b[o], b[o + 1], b[o + 2], b[o + 3]];
if s {
f32::from_be_bytes(v)
} else {
f32::from_le_bytes(v)
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CsiSample {
pub amplitude: Vec<f32>,
pub phase: Vec<f32>,
pub timestamp_ms: u64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BodyPartUV {
pub part_id: u8,
pub u_coords: Vec<f32>,
pub v_coords: Vec<f32>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PoseLabel {
pub keypoints: [(f32, f32, f32); 17],
pub body_parts: Vec<BodyPartUV>,
pub confidence: f32,
}
impl Default for PoseLabel {
fn default() -> Self {
Self {
keypoints: [(0.0, 0.0, 0.0); 17],
body_parts: Vec::new(),
confidence: 0.0,
}
}
}
pub struct SubcarrierResampler;
impl SubcarrierResampler {
pub fn resample(input: &[f32], from: usize, to: usize) -> Vec<f32> {
if from == to || from == 0 || to == 0 {
return input.to_vec();
}
if from < to {
Self::zero_pad(input, from, to)
} else {
Self::interpolate(input, from, to)
}
}
pub fn resample_phase(input: &[f32], from: usize, to: usize) -> Vec<f32> {
if from == to || from == 0 || to == 0 {
return input.to_vec();
}
let unwrapped = Self::phase_unwrap(input);
let resampled = if from < to {
Self::zero_pad(&unwrapped, from, to)
} else {
Self::interpolate(&unwrapped, from, to)
};
let pi = std::f32::consts::PI;
resampled
.iter()
.map(|&p| {
let mut w = p % (2.0 * pi);
if w > pi {
w -= 2.0 * pi;
}
if w < -pi {
w += 2.0 * pi;
}
w
})
.collect()
}
fn zero_pad(input: &[f32], from: usize, to: usize) -> Vec<f32> {
let pad_left = (to - from) / 2;
let mut out = vec![0.0f32; to];
for i in 0..from.min(input.len()) {
if pad_left + i < to {
out[pad_left + i] = input[i];
}
}
out
}
fn interpolate(input: &[f32], from: usize, to: usize) -> Vec<f32> {
let n = input.len().min(from);
if n <= 1 {
return vec![input.first().copied().unwrap_or(0.0); to];
}
(0..to)
.map(|i| {
let pos = i as f64 * (n - 1) as f64 / (to - 1).max(1) as f64;
let lo = pos.floor() as usize;
let hi = (lo + 1).min(n - 1);
let f = (pos - lo as f64) as f32;
input[lo] * (1.0 - f) + input[hi] * f
})
.collect()
}
fn phase_unwrap(phase: &[f32]) -> Vec<f32> {
let pi = std::f32::consts::PI;
let mut out = vec![0.0f32; phase.len()];
if phase.is_empty() {
return out;
}
out[0] = phase[0];
for i in 1..phase.len() {
let mut d = phase[i] - phase[i - 1];
while d > pi {
d -= 2.0 * pi;
}
while d < -pi {
d += 2.0 * pi;
}
out[i] = out[i - 1] + d;
}
out
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MmFiDataset {
pub csi_frames: Vec<CsiSample>,
pub labels: Vec<PoseLabel>,
pub sample_rate_hz: f32,
pub n_subcarriers: usize,
}
impl MmFiDataset {
pub const SUBCARRIERS: usize = 56;
pub fn load_from_directory(path: &Path) -> Result<Self> {
if !path.is_dir() {
return Err(DatasetError::Missing(format!(
"directory not found: {}",
path.display()
)));
}
let amp = NpyReader::read_file(&Self::find(path, &["csi_amplitude.npy", "csi.npy"])?)?;
let n = amp.shape.first().copied().unwrap_or(0);
let raw_sc = if amp.shape.len() >= 2 {
amp.shape[1]
} else {
amp.data.len() / n.max(1)
};
let phase_arr = Self::find(path, &["csi_phase.npy"])
.ok()
.and_then(|p| NpyReader::read_file(&p).ok());
let lab = NpyReader::read_file(&Self::find(path, &["labels.npy", "keypoints.npy"])?)?;
let mut csi_frames = Vec::with_capacity(n);
let mut labels = Vec::with_capacity(n);
for i in 0..n {
let s = i * raw_sc;
if s + raw_sc > amp.data.len() {
break;
}
let amplitude =
SubcarrierResampler::resample(&.data[s..s + raw_sc], raw_sc, Self::SUBCARRIERS);
let phase = phase_arr
.as_ref()
.map(|pa| {
let ps = i * raw_sc;
if ps + raw_sc <= pa.data.len() {
SubcarrierResampler::resample_phase(
&pa.data[ps..ps + raw_sc],
raw_sc,
Self::SUBCARRIERS,
)
} else {
vec![0.0; Self::SUBCARRIERS]
}
})
.unwrap_or_else(|| vec![0.0; Self::SUBCARRIERS]);
csi_frames.push(CsiSample {
amplitude,
phase,
timestamp_ms: i as u64 * 50,
});
let ks = i * 17 * 3;
let label = if ks + 51 <= lab.data.len() {
let d = &lab.data[ks..ks + 51];
let mut kp = [(0.0f32, 0.0, 0.0); 17];
for k in 0..17 {
kp[k] = (d[k * 3], d[k * 3 + 1], d[k * 3 + 2]);
}
PoseLabel {
keypoints: kp,
body_parts: Vec::new(),
confidence: 1.0,
}
} else {
PoseLabel::default()
};
labels.push(label);
}
Ok(Self {
csi_frames,
labels,
sample_rate_hz: 20.0,
n_subcarriers: Self::SUBCARRIERS,
})
}
pub fn resample_subcarriers(&mut self, from: usize, to: usize) {
for f in &mut self.csi_frames {
f.amplitude = SubcarrierResampler::resample(&f.amplitude, from, to);
f.phase = SubcarrierResampler::resample_phase(&f.phase, from, to);
}
self.n_subcarriers = to;
}
pub fn iter_windows(
&self,
ws: usize,
stride: usize,
) -> impl Iterator<Item = (&[CsiSample], &[PoseLabel])> {
let stride = stride.max(1);
let n = self.csi_frames.len();
(0..n)
.step_by(stride)
.filter(move |&s| s + ws <= n)
.map(move |s| (&self.csi_frames[s..s + ws], &self.labels[s..s + ws]))
}
pub fn split_train_val(self, ratio: f32) -> (Self, Self) {
let split = (self.csi_frames.len() as f32 * ratio.clamp(0.0, 1.0)) as usize;
let (tc, vc) = self.csi_frames.split_at(split);
let (tl, vl) = self.labels.split_at(split);
let mk = |c: &[CsiSample], l: &[PoseLabel]| Self {
csi_frames: c.to_vec(),
labels: l.to_vec(),
sample_rate_hz: self.sample_rate_hz,
n_subcarriers: self.n_subcarriers,
};
(mk(tc, tl), mk(vc, vl))
}
pub fn len(&self) -> usize {
self.csi_frames.len()
}
pub fn is_empty(&self) -> bool {
self.csi_frames.is_empty()
}
pub fn get(&self, idx: usize) -> Option<(&CsiSample, &PoseLabel)> {
self.csi_frames.get(idx).zip(self.labels.get(idx))
}
fn find(dir: &Path, names: &[&str]) -> Result<PathBuf> {
for n in names {
let p = dir.join(n);
if p.exists() {
return Ok(p);
}
}
Err(DatasetError::Missing(format!(
"none of {names:?} in {}",
dir.display()
)))
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WiPoseDataset {
pub csi_frames: Vec<CsiSample>,
pub labels: Vec<PoseLabel>,
pub sample_rate_hz: f32,
pub n_subcarriers: usize,
}
impl WiPoseDataset {
pub const RAW_SUBCARRIERS: usize = 30;
pub const TARGET_SUBCARRIERS: usize = 56;
pub const RAW_KEYPOINTS: usize = 18;
pub const COCO_KEYPOINTS: usize = 17;
pub fn load_from_mat(path: &Path) -> Result<Self> {
let arrays = MatReader::read_file(path)?;
let csi = arrays
.get("csi")
.or_else(|| arrays.get("csi_data"))
.or_else(|| arrays.get("CSI"))
.ok_or_else(|| DatasetError::Missing("no CSI variable in .mat".into()))?;
let n = csi.shape.first().copied().unwrap_or(0);
let raw = if csi.shape.len() >= 2 {
csi.shape[1]
} else {
Self::RAW_SUBCARRIERS
};
let lab = arrays
.get("keypoints")
.or_else(|| arrays.get("labels"))
.or_else(|| arrays.get("pose"));
let mut csi_frames = Vec::with_capacity(n);
let mut labels = Vec::with_capacity(n);
for i in 0..n {
let s = i * raw;
if s + raw > csi.data.len() {
break;
}
let amp =
SubcarrierResampler::resample(&csi.data[s..s + raw], raw, Self::TARGET_SUBCARRIERS);
csi_frames.push(CsiSample {
amplitude: amp,
phase: vec![0.0; Self::TARGET_SUBCARRIERS],
timestamp_ms: i as u64 * 100,
});
let label = lab
.and_then(|la| {
let ks = i * Self::RAW_KEYPOINTS * 3;
if ks + Self::RAW_KEYPOINTS * 3 <= la.data.len() {
Some(Self::map_18_to_17(
&la.data[ks..ks + Self::RAW_KEYPOINTS * 3],
))
} else {
None
}
})
.unwrap_or_default();
labels.push(label);
}
Ok(Self {
csi_frames,
labels,
sample_rate_hz: 10.0,
n_subcarriers: Self::TARGET_SUBCARRIERS,
})
}
fn map_18_to_17(data: &[f32]) -> PoseLabel {
let mut kp = [(0.0f32, 0.0, 0.0); 17];
if data.len() >= 18 * 3 {
kp[0] = (data[0], data[1], data[2]);
#[allow(clippy::needless_range_loop)]
for i in 1..17 {
let s = (i + 1) * 3;
kp[i] = (data[s], data[s + 1], data[s + 2]);
}
}
PoseLabel {
keypoints: kp,
body_parts: Vec::new(),
confidence: 1.0,
}
}
pub fn len(&self) -> usize {
self.csi_frames.len()
}
pub fn is_empty(&self) -> bool {
self.csi_frames.is_empty()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DataSource {
MmFi(PathBuf),
WiPose(PathBuf),
Combined(Vec<DataSource>),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataConfig {
pub source: DataSource,
pub window_size: usize,
pub stride: usize,
pub target_subcarriers: usize,
pub normalize: bool,
}
impl Default for DataConfig {
fn default() -> Self {
Self {
source: DataSource::Combined(Vec::new()),
window_size: 10,
stride: 5,
target_subcarriers: 56,
normalize: true,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingSample {
pub csi_window: Vec<Vec<f32>>,
pub pose_label: PoseLabel,
pub source: &'static str,
}
pub struct DataPipeline {
config: DataConfig,
}
impl DataPipeline {
pub fn new(config: DataConfig) -> Self {
Self { config }
}
pub fn load(&self) -> Result<Vec<TrainingSample>> {
let mut out = Vec::new();
self.load_source(&self.config.source, &mut out)?;
if self.config.normalize && !out.is_empty() {
Self::normalize_samples(&mut out);
}
Ok(out)
}
fn load_source(&self, src: &DataSource, out: &mut Vec<TrainingSample>) -> Result<()> {
match src {
DataSource::MmFi(p) => {
let mut ds = MmFiDataset::load_from_directory(p)?;
if ds.n_subcarriers != self.config.target_subcarriers {
let f = ds.n_subcarriers;
ds.resample_subcarriers(f, self.config.target_subcarriers);
}
self.extract_windows(&ds.csi_frames, &ds.labels, "mmfi", out);
}
DataSource::WiPose(p) => {
let ds = WiPoseDataset::load_from_mat(p)?;
self.extract_windows(&ds.csi_frames, &ds.labels, "wipose", out);
}
DataSource::Combined(srcs) => {
for s in srcs {
self.load_source(s, out)?;
}
}
}
Ok(())
}
fn extract_windows(
&self,
frames: &[CsiSample],
labels: &[PoseLabel],
source: &'static str,
out: &mut Vec<TrainingSample>,
) {
let (ws, stride) = (self.config.window_size, self.config.stride.max(1));
let mut s = 0;
while s + ws <= frames.len() {
let window: Vec<Vec<f32>> = frames[s..s + ws]
.iter()
.map(|f| f.amplitude.clone())
.collect();
let label = labels.get(s + ws / 2).cloned().unwrap_or_default();
out.push(TrainingSample {
csi_window: window,
pose_label: label,
source,
});
s += stride;
}
}
fn normalize_samples(samples: &mut [TrainingSample]) {
let ns = samples
.first()
.and_then(|s| s.csi_window.first())
.map(|f| f.len())
.unwrap_or(0);
if ns == 0 {
return;
}
let (mut sum, mut sq) = (vec![0.0f64; ns], vec![0.0f64; ns]);
let mut cnt = 0u64;
for s in samples.iter() {
for f in &s.csi_window {
for (j, &v) in f.iter().enumerate().take(ns) {
let v = v as f64;
sum[j] += v;
sq[j] += v * v;
}
cnt += 1;
}
}
if cnt == 0 {
return;
}
let mean: Vec<f64> = sum.iter().map(|s| s / cnt as f64).collect();
let std: Vec<f64> = sq
.iter()
.zip(mean.iter())
.map(|(&s, &m)| (s / cnt as f64 - m * m).max(0.0).sqrt().max(1e-8))
.collect();
for s in samples.iter_mut() {
for f in &mut s.csi_window {
for (j, v) in f.iter_mut().enumerate().take(ns) {
*v = ((*v as f64 - mean[j]) / std[j]) as f32;
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_npy_f32(shape: &[usize], data: &[f32]) -> Vec<u8> {
let ss = if shape.len() == 1 {
format!("({},)", shape[0])
} else {
format!(
"({})",
shape
.iter()
.map(|d| d.to_string())
.collect::<Vec<_>>()
.join(", ")
)
};
let hdr = format!("{{'descr': '<f4', 'fortran_order': False, 'shape': {ss}, }}");
let total = 10 + hdr.len();
let padded = total.div_ceil(64) * 64;
let hl = padded - 10;
let mut buf = Vec::new();
buf.extend_from_slice(b"\x93NUMPY\x01\x00");
buf.extend_from_slice(&(hl as u16).to_le_bytes());
buf.extend_from_slice(hdr.as_bytes());
buf.resize(10 + hl, b' ');
for &v in data {
buf.extend_from_slice(&v.to_le_bytes());
}
buf
}
fn make_npy_f64(shape: &[usize], data: &[f64]) -> Vec<u8> {
let ss = if shape.len() == 1 {
format!("({},)", shape[0])
} else {
format!(
"({})",
shape
.iter()
.map(|d| d.to_string())
.collect::<Vec<_>>()
.join(", ")
)
};
let hdr = format!("{{'descr': '<f8', 'fortran_order': False, 'shape': {ss}, }}");
let total = 10 + hdr.len();
let padded = total.div_ceil(64) * 64;
let hl = padded - 10;
let mut buf = Vec::new();
buf.extend_from_slice(b"\x93NUMPY\x01\x00");
buf.extend_from_slice(&(hl as u16).to_le_bytes());
buf.extend_from_slice(hdr.as_bytes());
buf.resize(10 + hl, b' ');
for &v in data {
buf.extend_from_slice(&v.to_le_bytes());
}
buf
}
#[test]
fn npy_header_parse_1d() {
let buf = make_npy_f32(&[5], &[1.0, 2.0, 3.0, 4.0, 5.0]);
let arr = NpyReader::parse(&buf).unwrap();
assert_eq!(arr.shape, vec![5]);
assert_eq!(arr.ndim(), 1);
assert_eq!(arr.len(), 5);
assert!((arr.data[0] - 1.0).abs() < f32::EPSILON);
assert!((arr.data[4] - 5.0).abs() < f32::EPSILON);
}
#[test]
fn npy_header_parse_2d() {
let data: Vec<f32> = (0..12).map(|i| i as f32).collect();
let buf = make_npy_f32(&[3, 4], &data);
let arr = NpyReader::parse(&buf).unwrap();
assert_eq!(arr.shape, vec![3, 4]);
assert_eq!(arr.ndim(), 2);
assert_eq!(arr.len(), 12);
}
#[test]
fn npy_header_parse_3d() {
let data: Vec<f64> = (0..24).map(|i| i as f64 * 0.5).collect();
let buf = make_npy_f64(&[2, 3, 4], &data);
let arr = NpyReader::parse(&buf).unwrap();
assert_eq!(arr.shape, vec![2, 3, 4]);
assert_eq!(arr.ndim(), 3);
assert_eq!(arr.len(), 24);
assert!((arr.data[23] - 11.5).abs() < 1e-5);
}
#[test]
fn subcarrier_resample_passthrough() {
let input: Vec<f32> = (0..56).map(|i| i as f32).collect();
let output = SubcarrierResampler::resample(&input, 56, 56);
assert_eq!(output, input);
}
#[test]
fn subcarrier_resample_upsample() {
let input: Vec<f32> = (0..30).map(|i| (i + 1) as f32).collect();
let out = SubcarrierResampler::resample(&input, 30, 56);
assert_eq!(out.len(), 56);
for (i, &v) in out[..13].iter().enumerate() {
assert!(v.abs() < f32::EPSILON, "expected zero at {i}");
}
for (&ov, &iv) in out[13..43].iter().zip(input.iter()) {
assert!((ov - iv).abs() < f32::EPSILON);
}
#[allow(clippy::needless_range_loop)]
for i in 43..56 {
assert!(out[i].abs() < f32::EPSILON, "expected zero at {i}");
}
}
#[test]
fn subcarrier_resample_downsample() {
let input: Vec<f32> = (0..114).map(|i| i as f32).collect();
let out = SubcarrierResampler::resample(&input, 114, 56);
assert_eq!(out.len(), 56);
assert!((out[0]).abs() < f32::EPSILON);
assert!((out[55] - 113.0).abs() < 0.1);
for i in 1..56 {
assert!(out[i] >= out[i - 1], "not monotonic at {i}");
}
}
#[test]
fn subcarrier_resample_preserves_dc() {
let out = SubcarrierResampler::resample(&vec![42.0f32; 114], 114, 56);
assert_eq!(out.len(), 56);
for (i, &v) in out.iter().enumerate() {
assert!((v - 42.0).abs() < 1e-5, "DC not preserved at {i}: {v}");
}
}
#[test]
fn mmfi_sample_structure() {
let s = CsiSample {
amplitude: vec![0.0; 56],
phase: vec![0.0; 56],
timestamp_ms: 100,
};
assert_eq!(s.amplitude.len(), 56);
assert_eq!(s.phase.len(), 56);
}
#[test]
fn wipose_zero_pad() {
let raw: Vec<f32> = (1..=30).map(|i| i as f32).collect();
let p = SubcarrierResampler::resample(&raw, 30, 56);
assert_eq!(p.len(), 56);
assert!(p[0].abs() < f32::EPSILON);
assert!((p[13] - 1.0).abs() < f32::EPSILON);
assert!((p[42] - 30.0).abs() < f32::EPSILON);
assert!(p[55].abs() < f32::EPSILON);
}
#[test]
fn wipose_keypoint_mapping() {
let mut kp = vec![0.0f32; 18 * 3];
kp[0] = 1.0;
kp[1] = 2.0;
kp[2] = 1.0; kp[3] = 99.0;
kp[4] = 99.0;
kp[5] = 99.0; kp[6] = 3.0;
kp[7] = 4.0;
kp[8] = 1.0; let label = WiPoseDataset::map_18_to_17(&kp);
assert_eq!(label.keypoints.len(), 17);
assert!((label.keypoints[0].0 - 1.0).abs() < f32::EPSILON);
assert!((label.keypoints[1].0 - 3.0).abs() < f32::EPSILON); }
#[test]
fn train_val_split_ratio() {
let mk = |n: usize| MmFiDataset {
csi_frames: (0..n)
.map(|i| CsiSample {
amplitude: vec![i as f32; 56],
phase: vec![0.0; 56],
timestamp_ms: i as u64,
})
.collect(),
labels: (0..n).map(|_| PoseLabel::default()).collect(),
sample_rate_hz: 20.0,
n_subcarriers: 56,
};
let (train, val) = mk(100).split_train_val(0.8);
assert_eq!(train.len(), 80);
assert_eq!(val.len(), 20);
assert_eq!(train.len() + val.len(), 100);
}
#[test]
fn sliding_window_count() {
let ds = MmFiDataset {
csi_frames: (0..20)
.map(|i| CsiSample {
amplitude: vec![i as f32; 56],
phase: vec![0.0; 56],
timestamp_ms: i as u64,
})
.collect(),
labels: (0..20).map(|_| PoseLabel::default()).collect(),
sample_rate_hz: 20.0,
n_subcarriers: 56,
};
assert_eq!(ds.iter_windows(5, 5).count(), 4);
assert_eq!(ds.iter_windows(5, 1).count(), 16);
}
#[test]
fn sliding_window_overlap() {
let ds = MmFiDataset {
csi_frames: (0..10)
.map(|i| CsiSample {
amplitude: vec![i as f32; 56],
phase: vec![0.0; 56],
timestamp_ms: i as u64,
})
.collect(),
labels: (0..10).map(|_| PoseLabel::default()).collect(),
sample_rate_hz: 20.0,
n_subcarriers: 56,
};
let w: Vec<_> = ds.iter_windows(4, 2).collect();
assert_eq!(w.len(), 4);
assert!((w[0].0[0].amplitude[0]).abs() < f32::EPSILON);
assert!((w[1].0[0].amplitude[0] - 2.0).abs() < f32::EPSILON);
assert_eq!(w[0].0[2].amplitude[0], w[1].0[0].amplitude[0]); }
#[test]
fn data_pipeline_normalize() {
let mut samples = vec![
TrainingSample {
csi_window: vec![vec![10.0, 20.0, 30.0]; 2],
pose_label: PoseLabel::default(),
source: "test",
},
TrainingSample {
csi_window: vec![vec![30.0, 40.0, 50.0]; 2],
pose_label: PoseLabel::default(),
source: "test",
},
];
DataPipeline::normalize_samples(&mut samples);
for j in 0..3 {
let (mut s, mut c) = (0.0f64, 0u64);
for sam in &samples {
for f in &sam.csi_window {
s += f[j] as f64;
c += 1;
}
}
assert!((s / c as f64).abs() < 1e-5, "mean not ~0 for sub {j}");
let mut vs = 0.0f64;
let m = s / c as f64;
for sam in &samples {
for f in &sam.csi_window {
vs += (f[j] as f64 - m).powi(2);
}
}
assert!(
((vs / c as f64).sqrt() - 1.0).abs() < 0.1,
"std not ~1 for sub {j}"
);
}
}
#[test]
fn pose_label_default() {
let l = PoseLabel::default();
assert_eq!(l.keypoints.len(), 17);
assert!(l.body_parts.is_empty());
assert!(l.confidence.abs() < f32::EPSILON);
for (i, kp) in l.keypoints.iter().enumerate() {
assert!(
kp.0.abs() < f32::EPSILON && kp.1.abs() < f32::EPSILON,
"kp {i} not zero"
);
}
}
#[test]
fn body_part_uv_round_trip() {
let bpu = BodyPartUV {
part_id: 5,
u_coords: vec![0.1, 0.2, 0.3],
v_coords: vec![0.4, 0.5, 0.6],
};
let json = serde_json::to_string(&bpu).unwrap();
let r: BodyPartUV = serde_json::from_str(&json).unwrap();
assert_eq!(r.part_id, 5);
assert_eq!(r.u_coords.len(), 3);
assert!((r.u_coords[0] - 0.1).abs() < f32::EPSILON);
assert!((r.v_coords[2] - 0.6).abs() < f32::EPSILON);
}
#[test]
fn combined_source_merges_datasets() {
let mk = |n: usize, base: f32| -> (Vec<CsiSample>, Vec<PoseLabel>) {
let f: Vec<CsiSample> = (0..n)
.map(|i| CsiSample {
amplitude: vec![base + i as f32; 56],
phase: vec![0.0; 56],
timestamp_ms: i as u64 * 50,
})
.collect();
let l: Vec<PoseLabel> = (0..n).map(|_| PoseLabel::default()).collect();
(f, l)
};
let pipe = DataPipeline::new(DataConfig {
source: DataSource::Combined(Vec::new()),
window_size: 3,
stride: 1,
target_subcarriers: 56,
normalize: false,
});
let mut all = Vec::new();
let (fa, la) = mk(5, 0.0);
pipe.extract_windows(&fa, &la, "mmfi", &mut all);
assert_eq!(all.len(), 3);
let (fb, lb) = mk(4, 100.0);
pipe.extract_windows(&fb, &lb, "wipose", &mut all);
assert_eq!(all.len(), 5);
assert_eq!(all[0].source, "mmfi");
assert_eq!(all[3].source, "wipose");
assert!(all[0].csi_window[0][0] < 10.0);
assert!(all[4].csi_window[0][0] > 90.0);
}
}