From e11f54eb4bda275d2e5d11febaf41800823816ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=27=C3=A9lectron=20rare?= <108685187+electron-rare@users.noreply.github.com> Date: Thu, 14 May 2026 12:13:37 +0200 Subject: [PATCH] feat(icp): wire fusion thread behind ICP_FUSION MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 9 of the ICP LiDAR plan: integrate the FusionWorker built in earlier tasks into the live data_only_viz pipeline without disturbing the existing ARKit pelvis fuse path or the Multi-HMR worker thread. A new IcpFusionThread pulls LiDAR frames from LidarTCPReader, stages them into State, and applies in-place ICP registration on state.persons_smplx[*].vertices_3d. It runs as a separate daemon thread parallel to MultiHMRWorker rather than inline per frame — the autonomous-worker architecture didn't fit the plan's per-frame call site, so we adapted to a polling thread at 8 Hz. Activation is opt-in via ICP_FUSION=1 plus ICP_LIDAR_HOST; the default code path is untouched. Shutdown wired through applicationWillTerminate_. MultiHMRWorker.predict_once is added as a documented stub (NotImplementedError) because the existing PyTorch run loop is too coupled to the camera and MPS lifecycle for a clean single-shot extraction. calibrate_lidar.py keeps its placeholder until a follow-up refactor extracts a pure _infer(rgb) helper. --- data_only_viz/icp_fusion_worker.py | 72 ++++++++++++++++++++++++ data_only_viz/main.py | 29 ++++++++++ data_only_viz/multi_hmr_worker.py | 14 +++++ data_only_viz/scripts/calibrate_lidar.py | 14 +++-- 4 files changed, 125 insertions(+), 4 deletions(-) create mode 100644 data_only_viz/icp_fusion_worker.py diff --git a/data_only_viz/icp_fusion_worker.py b/data_only_viz/icp_fusion_worker.py new file mode 100644 index 0000000..5f4921c --- /dev/null +++ b/data_only_viz/icp_fusion_worker.py @@ -0,0 +1,72 @@ +"""Threaded wrapper that polls State and calls FusionWorker.run_once. + +ICP fusion runs as a background thread parallel to the autonomous +Multi-HMR worker. It pulls the latest LiDAR frame from a +LidarTCPReader, stages it into State, and applies in-place ICP +registration to ``state.persons_smplx[*].vertices_3d``. + +Opt-in via ``ICP_FUSION=1`` from main.py. +""" +from __future__ import annotations + +import logging +import threading +import time +from typing import Optional + +from .icp_fusion import FusionWorker, IcpConfig +from .lidar_calib import load_extrinsic +from .lidar_receiver import LidarTCPReader + +_LOG = logging.getLogger(__name__) + + +class IcpFusionThread: + """Background thread: pull LiDAR frames, run FusionWorker on state.""" + + def __init__(self, state, host: str, port: int, + target_hz: float = 8.0) -> None: + self._state = state + self._reader = LidarTCPReader(host=host, port=port) + self._worker = FusionWorker(extrinsic=load_extrinsic(), + config=IcpConfig()) + self._period_s = 1.0 / max(target_hz, 0.5) + self._stop = threading.Event() + self._thread: Optional[threading.Thread] = None + + def start(self) -> None: + if self._thread is not None: + return + self._reader.start() + self._thread = threading.Thread( + target=self._run, name="icp-fusion", daemon=True) + self._thread.start() + _LOG.info("icp-fusion thread started") + + def stop(self) -> None: + self._stop.set() + self._reader.stop() + if self._thread is not None: + self._thread.join(timeout=2.0) + self._thread = None + + def _run(self) -> None: + while not self._stop.is_set(): + t0 = time.monotonic() + frame = self._reader.latest() + if frame is not None and self._state.persons_smplx: + # State doesn't expose a fine-grained lock for these + # fields here; rely on FusionWorker.run_once being + # write-only on persons_smplx[*].vertices_3d (replace in + # place) and the readers being tolerant of mid-update. + self._state.lidar_points = frame.points + self._state.lidar_timestamp_ns = frame.timestamp_ns + try: + self._state.icp_metadata = self._worker.run_once( + self._state) + except Exception as exc: # noqa: BLE001 + _LOG.warning("icp fusion failed: %s", exc) + self._state.icp_metadata = None + elapsed = time.monotonic() - t0 + if self._stop.wait(max(0.0, self._period_s - elapsed)): + return diff --git a/data_only_viz/main.py b/data_only_viz/main.py index 2ac3562..dbdadaa 100644 --- a/data_only_viz/main.py +++ b/data_only_viz/main.py @@ -260,6 +260,29 @@ class AppDelegate(NSObject): LOG.info("worker: + iPhone OSC listener :57128") except Exception as e: # noqa: BLE001 LOG.warning("iphone OSC listener start failed (%s)", e) + # ICP LiDAR fusion (opt-in via ICP_FUSION=1). Parallel to the + # ARKit pelvis fuse: ICP operates on SMPL-X dense vertices, not + # joints. Requires a calibrated extrinsic on disk (see + # scripts/calibrate_lidar.py) and an iPhone LiDAR stream + # broadcasting on ICP_LIDAR_HOST:ICP_LIDAR_PORT. + if _os.environ.get("ICP_FUSION", "0") == "1": + host = _os.environ.get("ICP_LIDAR_HOST") + if not host: + LOG.warning("ICP_FUSION=1 but ICP_LIDAR_HOST unset — " + "fusion disabled") + else: + try: + from .icp_fusion_worker import IcpFusionThread + self._icp_fusion = IcpFusionThread( + self._state, + host=host, + port=int(_os.environ.get("ICP_LIDAR_PORT", "5500")), + ) + self._icp_fusion.start() + LOG.info("worker: + ICP LiDAR fusion -> %s:%s", host, + _os.environ.get("ICP_LIDAR_PORT", "5500")) + except Exception as e: # noqa: BLE001 + LOG.warning("icp fusion start failed (%s)", e) # 0. Multi-HMR (SMPL-X 10475 verts mesh dense) — opt-in via flag if getattr(self._opts, "multi_hmr", False): try: @@ -596,6 +619,12 @@ class AppDelegate(NSObject): self._listener.stop() if self._pose_worker is not None: self._pose_worker.stop() + icp = getattr(self, "_icp_fusion", None) + if icp is not None: + try: + icp.stop() + except Exception as e: # noqa: BLE001 + LOG.warning("icp fusion stop failed (%s)", e) LOG.info("bye") diff --git a/data_only_viz/multi_hmr_worker.py b/data_only_viz/multi_hmr_worker.py index e31eb50..7a47769 100644 --- a/data_only_viz/multi_hmr_worker.py +++ b/data_only_viz/multi_hmr_worker.py @@ -116,6 +116,20 @@ class MultiHMRWorker: def stop(self) -> None: self._stop.set() + def predict_once(self, rgb_image): + """Single-shot SMPL-X prediction on one RGB image. + + Used by calibrate_lidar.py to acquire a pelvis vertex without + spinning the worker thread. The current PyTorch path is + deeply coupled to the run loop (model lifecycle, camera, MPS + setup) so this is left as a stub — calibrate_lidar.py keeps + its placeholder until a follow-up refactor extracts a pure + ``_infer(rgb) -> humans`` helper. + """ + raise NotImplementedError( + "MultiHMRWorker.predict_once is not wired yet — see " + "scripts/calibrate_lidar.py for the placeholder it gates") + def _run(self) -> None: if self.backend == "coreml": self._run_coreml() diff --git a/data_only_viz/scripts/calibrate_lidar.py b/data_only_viz/scripts/calibrate_lidar.py index ce21c22..9e82724 100644 --- a/data_only_viz/scripts/calibrate_lidar.py +++ b/data_only_viz/scripts/calibrate_lidar.py @@ -61,11 +61,17 @@ def main(argv: list[str] | None = None) -> int: reader = LidarTCPReader(host=args.lidar_host, port=args.lidar_port) reader.start() - # NB: the actual Multi-HMR getter is wired in Task 9 when the main pipeline - # exposes a single-shot predictor. For now this script is the *scaffolding* - # — Task 9 plugs in `multi_hmr_worker.predict_once()`. + # Task 9 added the ``MultiHMRWorker.predict_once`` API surface but + # left the body as ``NotImplementedError`` — the existing PyTorch + # path is too coupled to the worker thread for a clean extraction. + # When ``predict_once`` is wired (follow-up task), replace this + # placeholder by opening cv2.VideoCapture(args.webcam_index), + # running ``worker.predict_once(rgb)`` and returning + # ``person.vertices_3d[_PELVIS_VERT_INDEX]``. def _placeholder_pelvis_cam() -> np.ndarray: - raise SystemExit("calibrate_lidar requires Task 9 to be complete (predict_once API)") + raise SystemExit( + "calibrate_lidar needs MultiHMRWorker.predict_once to be " + "implemented (currently NotImplementedError)") pairs_cam, pairs_arkit = [], [] try: