diff --git a/data_only_viz/icp_fusion.py b/data_only_viz/icp_fusion.py index b1ff14f..f668aca 100644 --- a/data_only_viz/icp_fusion.py +++ b/data_only_viz/icp_fusion.py @@ -158,3 +158,58 @@ def partition_lidar_by_pid( continue out[pid] = pts[sel] return out + + +PELVIS_VERT_INDEX = 5559 # SMPL-X canonical pelvis vertex + + +@dataclass +class FusionMetadata: + applied: set[int] + fitness: dict[int, float] + rmse_m: dict[int, float] + n_lidar_points_used: int + + +class FusionWorker: + """Per-frame ICP fusion orchestrator (caller-driven, no internal thread).""" + + def __init__(self, extrinsic, config: IcpConfig | None = None) -> None: + self._extrinsic = extrinsic + self._config = config or IcpConfig() + + def set_extrinsic(self, extrinsic) -> None: + self._extrinsic = extrinsic + + def run_once(self, state) -> FusionMetadata: + applied: set[int] = set() + fitness: dict[int, float] = {} + rmse: dict[int, float] = {} + + lidar = getattr(state, "lidar_points", None) + if lidar is None or getattr(lidar, "size", 0) == 0 or not state.persons_smplx: + return FusionMetadata(applied, fitness, rmse, 0) + + T = np.asarray(self._extrinsic.T_arkit_to_cam, dtype=np.float32) + homog = np.concatenate([lidar, np.ones((lidar.shape[0], 1), dtype=np.float32)], axis=1) + lidar_cam = (homog @ T.T)[:, :3] + + pelvises = { + p.pid: p.vertices_3d[PELVIS_VERT_INDEX] + for p in state.persons_smplx + if p.vertices_3d is not None + } + parts = partition_lidar_by_pid(lidar_cam, pelvises, max_dist_m=1.0) + + for person in state.persons_smplx: + pts = parts.get(person.pid) + if pts is None: + continue + result = register_mesh_to_lidar(person.vertices_3d, pts, self._config) + fitness[person.pid] = result.fitness + rmse[person.pid] = result.rmse_m + if result.accepted: + person.vertices_3d = result.vertices_registered + applied.add(person.pid) + + return FusionMetadata(applied, fitness, rmse, lidar_cam.shape[0]) diff --git a/data_only_viz/state.py b/data_only_viz/state.py index 21b99fc..af64d60 100644 --- a/data_only_viz/state.py +++ b/data_only_viz/state.py @@ -139,6 +139,14 @@ class State: persons_arkit_joints: dict = field(default_factory=dict) persons_arkit_last_t: dict = field(default_factory=dict) + # ---- LiDAR / ICP mesh fusion (Task 8 - 2026-05-14) ---- + # Set by the LidarTCPReader poller; consumed by FusionWorker.run_once. + # The mesh-level fusion is complementary to the ARKit *joint* fusion + # above: joints are sparse + 60 Hz, LiDAR is dense + 5-10 Hz. + lidar_points: object = None # np.ndarray (N, 3) float32 ARKit world; None if no frame + lidar_timestamp_ns: int = 0 + icp_metadata: object = None # FusionMetadata from icp_fusion or None + # v1.3: centralised webcam source. WebcamSource owns the single # cv2.VideoCapture on the host and writes BGR frames here so all # consumers (MediaPipe Multi, Apple Vision, Multi-HMR worker, diff --git a/data_only_viz/tests/test_icp_fusion.py b/data_only_viz/tests/test_icp_fusion.py index 0babf8e..ed48a00 100644 --- a/data_only_viz/tests/test_icp_fusion.py +++ b/data_only_viz/tests/test_icp_fusion.py @@ -96,3 +96,49 @@ def test_partition_returns_empty_dict_when_no_pelvises() -> None: out = partition_lidar_by_pid(np.zeros((100, 3), dtype=np.float32), pelvises={}, max_dist_m=1.0) assert out == {} + + +def test_fusion_worker_in_place_update(monkeypatch) -> None: + from data_only_viz.icp_fusion import FusionWorker, IcpConfig + from data_only_viz.lidar_calib import Extrinsic + from data_only_viz.state import SMPLXPerson, State + + src = _synthetic_smplx_torso(seed=30) + verts = np.zeros((10475, 3), dtype=np.float32) + verts[: src.shape[0]] = src + verts[5559] = src.mean(axis=0) + + person = SMPLXPerson(pid=0, vertices_3d=verts.copy()) + state = State() + state.persons_smplx = [person] + + lidar_pts = src + np.array([0.0, 0.04, 0.0], dtype=np.float32) + state.lidar_points = lidar_pts + state.lidar_timestamp_ns = 1 + + worker = FusionWorker( + extrinsic=Extrinsic.identity(), + config=IcpConfig(), + ) + metadata = worker.run_once(state) + + assert metadata.applied == {0} + delta = state.persons_smplx[0].vertices_3d[5559] - verts[5559] + assert 0.02 <= delta[1] <= 0.06 + + +def test_fusion_worker_skips_when_no_lidar() -> None: + from data_only_viz.icp_fusion import FusionWorker, IcpConfig + from data_only_viz.lidar_calib import Extrinsic + from data_only_viz.state import SMPLXPerson, State + + verts = np.zeros((10475, 3), dtype=np.float32) + verts[5559] = [0.0, 1.0, 2.0] + state = State() + state.persons_smplx = [SMPLXPerson(pid=0, vertices_3d=verts.copy())] + state.lidar_points = None + + worker = FusionWorker(extrinsic=Extrinsic.identity(), config=IcpConfig()) + metadata = worker.run_once(state) + assert metadata.applied == set() + np.testing.assert_array_equal(state.persons_smplx[0].vertices_3d, verts)