diff --git a/data_only_viz/lidar_receiver.py b/data_only_viz/lidar_receiver.py new file mode 100644 index 0000000..60b3fbd --- /dev/null +++ b/data_only_viz/lidar_receiver.py @@ -0,0 +1,44 @@ +"""TCP receiver for iPhone ARBodyTracker LiDAR ARMeshAnchor stream. + +Wire format (per frame, after the 4-byte big-endian length prefix consumed +by the socket reader): + + [uint64 BE timestamp_ns] + [uint32 BE vertex_count] + [float32 LE x y z] * vertex_count + +The decoder is pure and side-effect-free so it can be unit-tested without a +socket. The socket reader lives in a separate class (LidarTCPReader) so its +threading model is independently testable. +""" +from __future__ import annotations + +import struct +from dataclasses import dataclass + +import numpy as np + +_HEADER = struct.Struct(">QI") # timestamp_ns, vertex_count + + +@dataclass(frozen=True) +class LidarFrame: + """One decoded LiDAR frame from the iPhone.""" + + timestamp_ns: int + points: np.ndarray # shape (N, 3), float32, ARKit world frame (meters) + + +def decode_frame(body: bytes) -> LidarFrame: + """Decode a frame body (length prefix already stripped).""" + if len(body) < _HEADER.size: + raise ValueError(f"truncated frame: header needs {_HEADER.size} bytes, got {len(body)}") + timestamp_ns, vertex_count = _HEADER.unpack_from(body, 0) + if vertex_count == 0: + raise ValueError("vertex_count must be > 0") + expected = _HEADER.size + vertex_count * 12 + if len(body) < expected: + raise ValueError(f"truncated frame: need {expected} bytes for {vertex_count} verts, got {len(body)}") + raw = body[_HEADER.size : expected] + pts = np.frombuffer(raw, dtype=" bytes: + """Mimic the iPhone-side encoder for round-trip testing.""" + n = points.shape[0] + body = struct.pack(">Q", timestamp_ns) + struct.pack(">I", n) + points.astype("I", len(body)) + return header + body + + +def test_decode_lidar_frame_roundtrip() -> None: + from data_only_viz.lidar_receiver import LidarFrame, decode_frame + + pts = np.array([[0.1, 0.2, 0.3], [-1.0, 2.0, 5.5]], dtype=np.float32) + payload = _encode_frame(pts, timestamp_ns=1_700_000_000_000_000_000) + + body = payload[4:] + frame = decode_frame(body) + + assert isinstance(frame, LidarFrame) + assert frame.timestamp_ns == 1_700_000_000_000_000_000 + np.testing.assert_allclose(frame.points, pts, atol=1e-6) + + +def test_decode_lidar_frame_rejects_truncated() -> None: + from data_only_viz.lidar_receiver import decode_frame + + pts = np.array([[1.0, 2.0, 3.0]], dtype=np.float32) + body = ( + struct.pack(">Q", 0) + + struct.pack(">I", 1) + + pts.astype(" None: + from data_only_viz.lidar_receiver import decode_frame + + body = struct.pack(">Q", 0) + struct.pack(">I", 0) + with pytest.raises(ValueError, match="vertex_count"): + decode_frame(body)