L'électron rare
216 lines
7.3 KiB
Python
216 lines
7.3 KiB
Python
"""ICP fusion between Multi-HMR SMPL-X meshes and iPhone LiDAR point clouds.
|
|
|
|
All operations happen in the **webcam camera frame** (meters, OpenCV
|
|
convention: +X right, +Y down, +Z forward). LiDAR points must be
|
|
pre-transformed via `Extrinsic.T_arkit_to_cam`.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from dataclasses import dataclass
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
import open3d as o3d
|
|
except ImportError: # pragma: no cover - exercised via skipif at import sites
|
|
o3d = None # type: ignore[assignment]
|
|
|
|
_LOG = logging.getLogger(__name__)
|
|
|
|
MIN_LIDAR_POINTS = 200
|
|
MIN_FITNESS = 0.30
|
|
MAX_RMSE_M = 0.05
|
|
CROP_MARGIN_M = 0.30
|
|
|
|
|
|
@dataclass
|
|
class IcpConfig:
|
|
voxel_size_m: float = 0.02
|
|
max_correspondence_m: float = 0.05
|
|
max_iterations: int = 30
|
|
|
|
|
|
@dataclass
|
|
class IcpResult:
|
|
vertices_registered: np.ndarray
|
|
accepted: bool
|
|
fitness: float
|
|
rmse_m: float
|
|
iterations: int
|
|
|
|
|
|
def register_mesh_to_lidar(
|
|
smplx_verts_cam: np.ndarray,
|
|
lidar_points_cam: np.ndarray,
|
|
config: IcpConfig | None = None,
|
|
) -> IcpResult:
|
|
"""Register SMPL-X verts onto a cropped LiDAR neighborhood."""
|
|
if o3d is None:
|
|
raise RuntimeError("open3d not installed — install with `uv sync --extra lidar`")
|
|
|
|
cfg = config or IcpConfig()
|
|
src = np.ascontiguousarray(smplx_verts_cam, dtype=np.float32)
|
|
|
|
if not np.isfinite(src).all():
|
|
_LOG.debug("ICP rejected: NaN/Inf in SMPL-X verts")
|
|
return IcpResult(src, False, 0.0, float("inf"), 0)
|
|
|
|
lidar = _crop_to_bbox(lidar_points_cam, src, margin_m=CROP_MARGIN_M)
|
|
if lidar.shape[0] < MIN_LIDAR_POINTS or not np.isfinite(lidar).all():
|
|
_LOG.debug("ICP rejected: insufficient LiDAR points (%d)", lidar.shape[0])
|
|
return IcpResult(src, False, 0.0, float("inf"), 0)
|
|
|
|
src_pcd = _to_pcd(src, cfg.voxel_size_m, estimate_normals=True)
|
|
tgt_pcd = _to_pcd(lidar, cfg.voxel_size_m, estimate_normals=True)
|
|
|
|
if len(src_pcd.points) < 10 or len(tgt_pcd.points) < 10:
|
|
return IcpResult(src, False, 0.0, float("inf"), 0)
|
|
|
|
criteria = o3d.pipelines.registration.ICPConvergenceCriteria(
|
|
max_iteration=cfg.max_iterations,
|
|
relative_fitness=1e-6,
|
|
relative_rmse=1e-6,
|
|
)
|
|
# Coarse-to-fine: a wide first pass handles translations larger than the
|
|
# final correspondence threshold, then the strict pass refines and gates.
|
|
coarse = o3d.pipelines.registration.registration_icp(
|
|
src_pcd, tgt_pcd, max(cfg.max_correspondence_m * 5.0, 0.20),
|
|
np.eye(4),
|
|
o3d.pipelines.registration.TransformationEstimationPointToPlane(),
|
|
criteria,
|
|
)
|
|
result = o3d.pipelines.registration.registration_icp(
|
|
src_pcd, tgt_pcd, cfg.max_correspondence_m,
|
|
coarse.transformation,
|
|
o3d.pipelines.registration.TransformationEstimationPointToPlane(),
|
|
criteria,
|
|
)
|
|
|
|
accepted = (result.fitness >= MIN_FITNESS) and (result.inlier_rmse <= MAX_RMSE_M)
|
|
if not accepted:
|
|
_LOG.debug("ICP rejected: fitness=%.3f rmse=%.4f", result.fitness, result.inlier_rmse)
|
|
return IcpResult(src, False, float(result.fitness), float(result.inlier_rmse), 0)
|
|
|
|
T = np.asarray(result.transformation, dtype=np.float32)
|
|
homog = np.concatenate([src, np.ones((src.shape[0], 1), dtype=np.float32)], axis=1)
|
|
fused = (homog @ T.T)[:, :3]
|
|
if not np.isfinite(fused).all():
|
|
return IcpResult(src, False, float(result.fitness), float(result.inlier_rmse), 0)
|
|
|
|
return IcpResult(
|
|
vertices_registered=np.ascontiguousarray(fused, dtype=np.float32),
|
|
accepted=True,
|
|
fitness=float(result.fitness),
|
|
rmse_m=float(result.inlier_rmse),
|
|
iterations=cfg.max_iterations,
|
|
)
|
|
|
|
|
|
def _crop_to_bbox(points: np.ndarray, anchor: np.ndarray, margin_m: float) -> np.ndarray:
|
|
if points.size == 0:
|
|
return points.astype(np.float32, copy=False)
|
|
lo = anchor.min(axis=0) - margin_m
|
|
hi = anchor.max(axis=0) + margin_m
|
|
mask = np.all((points >= lo) & (points <= hi), axis=1)
|
|
return points[mask].astype(np.float32, copy=False)
|
|
|
|
|
|
def _to_pcd(points: np.ndarray, voxel_size_m: float, estimate_normals: bool):
|
|
pcd = o3d.geometry.PointCloud()
|
|
pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64, copy=False))
|
|
if voxel_size_m > 0:
|
|
pcd = pcd.voxel_down_sample(voxel_size_m)
|
|
if estimate_normals:
|
|
pcd.estimate_normals(
|
|
search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size_m * 2, max_nn=30),
|
|
)
|
|
return pcd
|
|
|
|
|
|
def partition_lidar_by_pid(
|
|
lidar_points_cam: np.ndarray,
|
|
pelvises: dict[int, np.ndarray],
|
|
max_dist_m: float = 1.0,
|
|
) -> dict[int, np.ndarray]:
|
|
"""Assign each LiDAR point to the closest pelvis within ``max_dist_m``.
|
|
|
|
Points beyond ``max_dist_m`` from every pelvis (background, furniture)
|
|
are dropped. Returns ``{pid: (M, 3) float32}`` — pids with zero assigned
|
|
points are omitted.
|
|
"""
|
|
if not pelvises or lidar_points_cam.size == 0:
|
|
return {}
|
|
pids = list(pelvises.keys())
|
|
centers = np.stack([pelvises[p] for p in pids]).astype(np.float32)
|
|
pts = np.ascontiguousarray(lidar_points_cam, dtype=np.float32)
|
|
|
|
diff = pts[:, None, :] - centers[None, :, :]
|
|
d2 = np.einsum("npk,npk->np", diff, diff)
|
|
nearest = d2.argmin(axis=1)
|
|
nearest_d = np.sqrt(d2[np.arange(d2.shape[0]), nearest])
|
|
|
|
mask = nearest_d <= max_dist_m
|
|
out: dict[int, np.ndarray] = {}
|
|
for idx, pid in enumerate(pids):
|
|
sel = mask & (nearest == idx)
|
|
if not sel.any():
|
|
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])
|