feat(av-live): openpos 3D + DINO reid + filter
Three improvements wired end-to-end: 1. Openpos 3D skeleton visible: Skeleton3DRenderer attached to a RealityKit AnchorEntity in BodyView, toggled by showSkeleton or vizMode==9. PoseOSCListener now parses /pose3d/count and /pose3d/kp (plus restored /face/* and /hand/* paths). 2. DINO re-id (dinov2_vits14, ~9 ms ANE forward): MeshRigger combines Hungarian IoU with cosine similarity over a per-pid embedding history (deque maxlen=10), weighted by MULTIHMR_REID_ALPHA (default 0.5). Falls back to pure IoU if DINO mlpackage absent or scipy missing. state.last_frame_rgb buffer added so the rigger can crop bbox regions for embedding. 3. PoseFilterChain on pose_world_landmarks: median (anti-spike) -> Kalman constant-velocity -> 50 ms lookahead -> IK elbow/knee/ankle clamp. Configurable via POSE_FILTER env (default median+kalman+lookahead+ik). <2 ms per frame for typical 1-2 persons. Tests: 5 new in test_dino_reid.py + 6 new in test_pose_filter.py, all green. Live validated by user: skeleton spawns, mesh stays stable.
This commit is contained in:
@@ -0,0 +1,204 @@
|
||||
"""DINOv2 ViT-S/14 person re-id backend (CoreML via pyobjc).
|
||||
|
||||
Loads the .mlpackage produced by ``scripts/convert_dinov2.py`` and runs
|
||||
inference one crop at a time (pyobjc + MLDictionaryFeatureProvider).
|
||||
Same pattern as ``multihmr_coreml.py`` so Python 3.14 works (no
|
||||
coremltools dependency at runtime).
|
||||
|
||||
Embeddings are L2-normalised inside the CoreML graph, so cosine sim
|
||||
between two outputs is a plain dot product.
|
||||
|
||||
Public API::
|
||||
|
||||
reid = DinoReid(mlpackage_path) # optional path
|
||||
emb = reid.embed_crops(list_of_uint8_HWC) # -> np.ndarray (N, 384)
|
||||
DinoReid.is_available() # bool
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
LOG = logging.getLogger("dino_reid")
|
||||
|
||||
DEFAULT_MLPACKAGE = (
|
||||
Path.home() / ".cache" / "av-live-multihmr" / "dinov2_vits14.mlpackage"
|
||||
)
|
||||
|
||||
EMBED_DIM = 384
|
||||
INPUT_SIZE = 224
|
||||
|
||||
# MLMultiArrayDataType raw values (from CoreML headers).
|
||||
ML_DTYPE_FLOAT32 = 65568
|
||||
ML_DTYPE_FLOAT16 = 65552
|
||||
ML_DTYPE_DOUBLE = 65600
|
||||
|
||||
|
||||
def _resize_crop(crop_uint8: np.ndarray) -> np.ndarray:
|
||||
"""Resize an HxWx3 uint8 crop to (3, 224, 224) float32 in [0, 1].
|
||||
|
||||
Uses ``cv2.resize`` when available, falls back to a simple stride
|
||||
sampler otherwise (avoids hard cv2 dep in test envs)."""
|
||||
if crop_uint8.ndim != 3 or crop_uint8.shape[2] != 3:
|
||||
raise ValueError(f"crop must be HxWx3 uint8, got {crop_uint8.shape}")
|
||||
if crop_uint8.shape[0] == INPUT_SIZE and crop_uint8.shape[1] == INPUT_SIZE:
|
||||
rgb = crop_uint8
|
||||
else:
|
||||
try:
|
||||
import cv2
|
||||
rgb = cv2.resize(crop_uint8, (INPUT_SIZE, INPUT_SIZE),
|
||||
interpolation=cv2.INTER_AREA)
|
||||
except ImportError:
|
||||
h, w = crop_uint8.shape[:2]
|
||||
ys = (np.linspace(0, h - 1, INPUT_SIZE)).astype(np.int32)
|
||||
xs = (np.linspace(0, w - 1, INPUT_SIZE)).astype(np.int32)
|
||||
rgb = crop_uint8[ys][:, xs]
|
||||
return (rgb.astype(np.float32) / 255.0).transpose(2, 0, 1)
|
||||
|
||||
|
||||
class DinoReid:
|
||||
"""Forward DINOv2 ViT-S/14 over RGB crops, return L2-normalised
|
||||
embeddings (N, 384)."""
|
||||
|
||||
def __init__(self, mlpackage_path: Path | str | None = None) -> None:
|
||||
self.path = Path(mlpackage_path) if mlpackage_path else DEFAULT_MLPACKAGE
|
||||
if not self.path.exists():
|
||||
raise FileNotFoundError(f"mlpackage missing: {self.path}")
|
||||
|
||||
import objc
|
||||
from Foundation import NSURL
|
||||
|
||||
self._objc = objc
|
||||
self._NSURL = NSURL
|
||||
|
||||
ns: dict = {}
|
||||
objc.loadBundle("CoreML", ns,
|
||||
"/System/Library/Frameworks/CoreML.framework")
|
||||
self._ns = ns
|
||||
|
||||
MLModel = ns["MLModel"]
|
||||
MLModelConfiguration = ns["MLModelConfiguration"]
|
||||
cfg = MLModelConfiguration.alloc().init()
|
||||
try:
|
||||
# 2 = MLComputeUnitsAll (CPU+GPU+ANE). DINOv2 ViT-S/14
|
||||
# converts cleanly and ANE serves it well.
|
||||
cfg.setComputeUnits_(2)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
url = NSURL.fileURLWithPath_(str(self.path))
|
||||
compiled = MLModel.compileModelAtURL_error_(url, None)
|
||||
if compiled is None:
|
||||
raise RuntimeError(f"compile failed for {self.path}")
|
||||
model = MLModel.modelWithContentsOfURL_configuration_error_(
|
||||
compiled, cfg, None)
|
||||
if model is None:
|
||||
raise RuntimeError(f"load failed for {compiled}")
|
||||
self._model = model
|
||||
|
||||
# Discover the output feature name (single tensor).
|
||||
desc = model.modelDescription()
|
||||
out_names = [str(n) for n in desc.outputDescriptionsByName().keys()]
|
||||
self._out_name = out_names[0] if out_names else "embedding"
|
||||
LOG.info("dino_reid loaded (%s, out=%s)", self.path.name,
|
||||
self._out_name)
|
||||
|
||||
@classmethod
|
||||
def is_available(cls, mlpackage_path: Path | str | None = None) -> bool:
|
||||
p = Path(mlpackage_path) if mlpackage_path else DEFAULT_MLPACKAGE
|
||||
if not p.exists():
|
||||
return False
|
||||
try:
|
||||
import objc # noqa: F401
|
||||
from Foundation import NSURL # noqa: F401
|
||||
return True
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# MLMultiArray plumbing — mirrors multihmr_coreml._np_to_mlarray /
|
||||
# _mlarray_to_np. Float32 in, float32-or-float16 out.
|
||||
# ------------------------------------------------------------------
|
||||
def _np_to_mlarray(self, arr: np.ndarray):
|
||||
import ctypes
|
||||
MLMultiArray = self._ns["MLMultiArray"]
|
||||
arr = np.ascontiguousarray(arr, dtype=np.float32)
|
||||
shape = [int(s) for s in arr.shape]
|
||||
ml = MLMultiArray.alloc().initWithShape_dataType_error_(
|
||||
shape, ML_DTYPE_FLOAT32, None)
|
||||
if ml is None:
|
||||
raise RuntimeError("MLMultiArray alloc failed")
|
||||
ptr = ml.dataPointer()
|
||||
addr = int(ptr) if isinstance(ptr, int) else ctypes.cast(
|
||||
ptr, ctypes.c_void_p).value
|
||||
if addr is None:
|
||||
raise RuntimeError("dataPointer null")
|
||||
ctypes.memmove(addr, arr.ctypes.data, arr.nbytes)
|
||||
return ml
|
||||
|
||||
def _mlarray_to_np(self, ml) -> np.ndarray:
|
||||
import ctypes
|
||||
shape = tuple(int(s) for s in ml.shape())
|
||||
dtype_id = int(ml.dataType())
|
||||
count = 1
|
||||
for s in shape:
|
||||
count *= s
|
||||
ptr = ml.dataPointer()
|
||||
addr = int(ptr) if isinstance(ptr, int) else ctypes.cast(
|
||||
ptr, ctypes.c_void_p).value
|
||||
if addr is None:
|
||||
raise RuntimeError("dataPointer null")
|
||||
if dtype_id == ML_DTYPE_FLOAT16:
|
||||
raw = (ctypes.c_uint16 * count).from_address(addr)
|
||||
arr = np.ctypeslib.as_array(raw).view(np.float16).astype(np.float32)
|
||||
elif dtype_id == ML_DTYPE_FLOAT32:
|
||||
raw = (ctypes.c_float * count).from_address(addr)
|
||||
arr = np.ctypeslib.as_array(raw).copy()
|
||||
elif dtype_id == ML_DTYPE_DOUBLE:
|
||||
raw = (ctypes.c_double * count).from_address(addr)
|
||||
arr = np.ctypeslib.as_array(raw).astype(np.float32)
|
||||
else:
|
||||
raise RuntimeError(f"unsupported dtype {dtype_id}")
|
||||
return arr.reshape(shape)
|
||||
|
||||
def _predict_one(self, image_chw: np.ndarray) -> np.ndarray:
|
||||
MLDictionaryFeatureProvider = self._ns["MLDictionaryFeatureProvider"]
|
||||
MLFeatureValue = self._ns["MLFeatureValue"]
|
||||
x4 = image_chw[np.newaxis, ...] if image_chw.ndim == 3 else image_chw
|
||||
img_ml = self._np_to_mlarray(x4)
|
||||
feats = {"image": MLFeatureValue.featureValueWithMultiArray_(img_ml)}
|
||||
provider = MLDictionaryFeatureProvider.alloc(
|
||||
).initWithDictionary_error_(feats, None)
|
||||
if provider is None:
|
||||
raise RuntimeError("provider alloc failed")
|
||||
out = self._model.predictionFromFeatures_error_(provider, None)
|
||||
if out is None:
|
||||
raise RuntimeError("predict failed")
|
||||
fv = out.featureValueForName_(self._out_name)
|
||||
ml = fv.multiArrayValue()
|
||||
return self._mlarray_to_np(ml).reshape(-1)
|
||||
|
||||
def embed_crops(
|
||||
self, crops_uint8: Sequence[np.ndarray],
|
||||
) -> np.ndarray:
|
||||
"""Embed a list of HxWx3 uint8 RGB crops -> (N, 384) float32.
|
||||
|
||||
Loops one crop at a time (the CoreML model is traced for B=1).
|
||||
For typical N <= 4 this is still 10-15 ms total on M5."""
|
||||
if not crops_uint8:
|
||||
return np.zeros((0, EMBED_DIM), dtype=np.float32)
|
||||
t0 = time.perf_counter()
|
||||
out = np.zeros((len(crops_uint8), EMBED_DIM), dtype=np.float32)
|
||||
for i, c in enumerate(crops_uint8):
|
||||
chw = _resize_crop(c)
|
||||
out[i] = self._predict_one(chw)
|
||||
dt_ms = (time.perf_counter() - t0) * 1e3
|
||||
if LOG.isEnabledFor(logging.DEBUG) or dt_ms > 50.0:
|
||||
LOG.log(
|
||||
logging.DEBUG if dt_ms <= 50.0 else logging.INFO,
|
||||
"embedded %d crops in %.1f ms", len(crops_uint8), dt_ms)
|
||||
return out
|
||||
@@ -14,6 +14,8 @@ Limitations connues (premiere iteration) :
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import logging
|
||||
import math
|
||||
import threading
|
||||
import time
|
||||
@@ -21,8 +23,16 @@ from dataclasses import dataclass, field
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
from scipy.optimize import linear_sum_assignment
|
||||
_HAVE_SCIPY = True
|
||||
except ImportError: # noqa: BLE001
|
||||
_HAVE_SCIPY = False
|
||||
|
||||
from .state import PoseKp, SMPLXPerson, State
|
||||
|
||||
LOG = logging.getLogger("mesh_rigger")
|
||||
|
||||
|
||||
# Indices MediaPipe POSE_LANDMARKS pour les hanches (pelvis 2D = midpoint).
|
||||
_LEFT_HIP = 23
|
||||
@@ -55,6 +65,70 @@ def _pelvis_2d_from_body(body: list[PoseKp]) -> tuple[float, float] | None:
|
||||
return (0.5 * (lh.x + rh.x), 0.5 * (lh.y + rh.y))
|
||||
|
||||
|
||||
def _body_bbox_norm(
|
||||
body: list[PoseKp],
|
||||
) -> tuple[float, float, float, float] | None:
|
||||
"""Bbox image-normalized [0,1] from a list of body landmarks
|
||||
(Vision 19 joints OR MediaPipe 33). None if not enough confident
|
||||
points."""
|
||||
if not body:
|
||||
return None
|
||||
xs = [kp.x for kp in body if kp.c > 0.05]
|
||||
ys = [kp.y for kp in body if kp.c > 0.05]
|
||||
if len(xs) < 4 or len(ys) < 4:
|
||||
return None
|
||||
x0, x1 = max(0.0, min(xs)), min(1.0, max(xs))
|
||||
y0, y1 = max(0.0, min(ys)), min(1.0, max(ys))
|
||||
# Pad 10% to capture full body silhouette.
|
||||
dx = (x1 - x0) * 0.10
|
||||
dy = (y1 - y0) * 0.10
|
||||
x0 = max(0.0, x0 - dx); x1 = min(1.0, x1 + dx)
|
||||
y0 = max(0.0, y0 - dy); y1 = min(1.0, y1 + dy)
|
||||
if x1 - x0 < 0.02 or y1 - y0 < 0.02:
|
||||
return None
|
||||
return (x0, y0, x1, y1)
|
||||
|
||||
|
||||
def _mesh_bbox_norm(p: SMPLXPerson) -> tuple[float, float, float, float] | None:
|
||||
"""Project SMPL-X mesh vertices to image-normalized bbox.
|
||||
|
||||
Multi-HMR uses focal = IMG_SIZE camera intrinsics. World verts
|
||||
have z>0 (in front of camera)."""
|
||||
v = np.asarray(p.vertices_3d, dtype=np.float32)
|
||||
if v.size == 0 or v.shape[0] < 100:
|
||||
return None
|
||||
z = v[:, 2]
|
||||
valid = z > 1e-3
|
||||
if not np.any(valid):
|
||||
return None
|
||||
x_img = (v[valid, 0] * _FOCAL / z[valid]) / _IMG_SIZE + 0.5
|
||||
y_img = (v[valid, 1] * _FOCAL / z[valid]) / _IMG_SIZE + 0.5
|
||||
x0, x1 = float(x_img.min()), float(x_img.max())
|
||||
y0, y1 = float(y_img.min()), float(y_img.max())
|
||||
x0 = max(0.0, x0); x1 = min(1.0, x1)
|
||||
y0 = max(0.0, y0); y1 = min(1.0, y1)
|
||||
if x1 - x0 < 0.02 or y1 - y0 < 0.02:
|
||||
return None
|
||||
return (x0, y0, x1, y1)
|
||||
|
||||
|
||||
def _iou_norm(
|
||||
a: tuple[float, float, float, float],
|
||||
b: tuple[float, float, float, float],
|
||||
) -> float:
|
||||
ax0, ay0, ax1, ay1 = a
|
||||
bx0, by0, bx1, by1 = b
|
||||
ix0 = max(ax0, bx0); iy0 = max(ay0, by0)
|
||||
ix1 = min(ax1, bx1); iy1 = min(ay1, by1)
|
||||
iw = max(0.0, ix1 - ix0); ih = max(0.0, iy1 - iy0)
|
||||
inter = iw * ih
|
||||
if inter <= 0:
|
||||
return 0.0
|
||||
a_area = (ax1 - ax0) * (ay1 - ay0)
|
||||
b_area = (bx1 - bx0) * (by1 - by0)
|
||||
return float(inter / (a_area + b_area - inter + 1e-9))
|
||||
|
||||
|
||||
def _vision_pid_match(
|
||||
keyframe_pelvis_2d: tuple[float, float] | None,
|
||||
vision_bodies: list[list[PoseKp]],
|
||||
@@ -89,14 +163,22 @@ class MeshRigger:
|
||||
Thread-safe : ne mute pas le state, retourne une nouvelle liste.
|
||||
"""
|
||||
|
||||
def __init__(self, state: State, hold_window_s: float = 1.5) -> None:
|
||||
def __init__(self, state: State, hold_window_s: float = 1.5,
|
||||
dino_weight: float = 0.5,
|
||||
dino_reid=None) -> None:
|
||||
self.state = state
|
||||
self.hold_window_s = hold_window_s
|
||||
self.dino_weight = float(dino_weight)
|
||||
self.dino_reid = dino_reid
|
||||
self._lock = threading.Lock()
|
||||
# pid Multi-HMR -> keyframe
|
||||
self._keyframes: dict[int, _Keyframe] = {}
|
||||
# pid Multi-HMR -> pid Vision matched (sticky across frames)
|
||||
self._vision_pid_map: dict[int, int] = {}
|
||||
# pid Multi-HMR -> recent DINO embeddings (mean -> reid signature)
|
||||
self._pid_embeddings: dict[int, collections.deque] = {}
|
||||
# Cached log throttle
|
||||
self._next_dino_log = 0.0
|
||||
|
||||
def apply(
|
||||
self,
|
||||
@@ -114,6 +196,14 @@ class MeshRigger:
|
||||
if old_pid not in current_pids:
|
||||
self._keyframes.pop(old_pid, None)
|
||||
self._vision_pid_map.pop(old_pid, None)
|
||||
self._pid_embeddings.pop(old_pid, None)
|
||||
|
||||
# 2) DINO fusion: if a reid backend is wired, try Hungarian
|
||||
# over (mesh keyframe pids) x (Vision body pids) using
|
||||
# alpha*IoU + (1-alpha)*cosine. This only kicks in when a
|
||||
# keyframe is detected this call AND we have an RGB frame.
|
||||
self._dino_match(persons_smplx, persons_body,
|
||||
persons_body_ids)
|
||||
|
||||
out: list[SMPLXPerson] = []
|
||||
for person in persons_smplx:
|
||||
@@ -199,6 +289,136 @@ class MeshRigger:
|
||||
))
|
||||
return out
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# DINOv2 reid hooks
|
||||
# ------------------------------------------------------------------
|
||||
def _dino_match(
|
||||
self,
|
||||
persons_smplx: list[SMPLXPerson],
|
||||
persons_body: list[list[PoseKp]],
|
||||
persons_body_ids: list[int],
|
||||
) -> None:
|
||||
"""Update self._vision_pid_map and self._pid_embeddings by
|
||||
matching mesh pids against Vision pids on alpha*IoU +
|
||||
(1-alpha)*DINO cosine. No-op if any prerequisite missing.
|
||||
|
||||
Caller must hold self._lock."""
|
||||
if self.dino_reid is None or not _HAVE_SCIPY:
|
||||
return
|
||||
if not persons_smplx or not persons_body:
|
||||
return
|
||||
# Need at least one new keyframe to be worth running DINO.
|
||||
new_kf_pids: list[int] = []
|
||||
for p in persons_smplx:
|
||||
kf = self._keyframes.get(p.pid)
|
||||
if kf is None or not np.allclose(
|
||||
kf.translation, p.translation, atol=1e-4):
|
||||
new_kf_pids.append(int(p.pid))
|
||||
if not new_kf_pids:
|
||||
return
|
||||
|
||||
# Acquire current RGB frame (best effort, no double lock).
|
||||
frame = self.state.last_frame_rgb
|
||||
if frame is None:
|
||||
return
|
||||
H, W = frame.shape[:2]
|
||||
|
||||
# Build Vision bboxes (image-normalized) and pixel crops.
|
||||
v_bboxes_norm: list[tuple[float, float, float, float]] = []
|
||||
v_crops: list[np.ndarray] = []
|
||||
v_pids: list[int] = []
|
||||
for body, vpid in zip(persons_body, persons_body_ids):
|
||||
bb = _body_bbox_norm(body)
|
||||
if bb is None:
|
||||
continue
|
||||
x0, y0, x1, y1 = bb
|
||||
px0 = max(0, int(x0 * W))
|
||||
py0 = max(0, int(y0 * H))
|
||||
px1 = min(W, int(x1 * W))
|
||||
py1 = min(H, int(y1 * H))
|
||||
if px1 <= px0 + 4 or py1 <= py0 + 4:
|
||||
continue
|
||||
v_bboxes_norm.append(bb)
|
||||
v_crops.append(frame[py0:py1, px0:px1].copy())
|
||||
v_pids.append(int(vpid))
|
||||
|
||||
if not v_crops:
|
||||
return
|
||||
|
||||
# Build mesh bboxes (image-normalized) from world pelvis proj.
|
||||
m_bboxes_norm: list[tuple[float, float, float, float]] = []
|
||||
m_pids_keep: list[int] = []
|
||||
m_crops: list[np.ndarray] = []
|
||||
for p in persons_smplx:
|
||||
bb = _mesh_bbox_norm(p)
|
||||
if bb is None:
|
||||
continue
|
||||
m_bboxes_norm.append(bb)
|
||||
m_pids_keep.append(int(p.pid))
|
||||
x0, y0, x1, y1 = bb
|
||||
px0 = max(0, int(x0 * W))
|
||||
py0 = max(0, int(y0 * H))
|
||||
px1 = min(W, int(x1 * W))
|
||||
py1 = min(H, int(y1 * H))
|
||||
if px1 > px0 + 4 and py1 > py0 + 4:
|
||||
m_crops.append(frame[py0:py1, px0:px1].copy())
|
||||
else:
|
||||
m_crops.append(None) # type: ignore[arg-type]
|
||||
|
||||
if not m_bboxes_norm:
|
||||
return
|
||||
|
||||
# Embed Vision crops in one batch (still loops internally).
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
v_emb = self.dino_reid.embed_crops(v_crops)
|
||||
except Exception as e: # noqa: BLE001
|
||||
LOG.warning("dino_reid embed failed: %s", e)
|
||||
return
|
||||
|
||||
# Build cost matrix mesh x vision : 1 - (alpha*IoU + (1-alpha)*cos)
|
||||
n_m = len(m_bboxes_norm)
|
||||
n_v = len(v_bboxes_norm)
|
||||
alpha = float(np.clip(self.dino_weight, 0.0, 1.0))
|
||||
cost = np.ones((n_m, n_v), dtype=np.float32)
|
||||
for i, mbb in enumerate(m_bboxes_norm):
|
||||
hist = self._pid_embeddings.get(m_pids_keep[i])
|
||||
mean_emb = None
|
||||
if hist:
|
||||
stack = np.stack(list(hist), axis=0)
|
||||
mean_emb = stack.mean(axis=0)
|
||||
n = np.linalg.norm(mean_emb) + 1e-8
|
||||
mean_emb = mean_emb / n
|
||||
for j, vbb in enumerate(v_bboxes_norm):
|
||||
iou = _iou_norm(mbb, vbb)
|
||||
if mean_emb is not None:
|
||||
cos = float(np.dot(mean_emb, v_emb[j]))
|
||||
else:
|
||||
cos = iou # no history -> trust IoU
|
||||
score = alpha * iou + (1.0 - alpha) * max(0.0, cos)
|
||||
cost[i, j] = 1.0 - score
|
||||
|
||||
rr, cc = linear_sum_assignment(cost)
|
||||
for i, j in zip(rr, cc):
|
||||
if cost[i, j] >= 0.95:
|
||||
continue # weak match, ignore
|
||||
mpid = m_pids_keep[i]
|
||||
self._vision_pid_map[mpid] = v_pids[j]
|
||||
# Update embedding history for THIS mesh pid using the
|
||||
# Vision crop (most recent visual evidence).
|
||||
dq = self._pid_embeddings.setdefault(
|
||||
mpid, collections.deque(maxlen=10))
|
||||
dq.append(v_emb[j].copy())
|
||||
|
||||
now = time.monotonic()
|
||||
dt_ms = (time.perf_counter() - t0) * 1e3
|
||||
if now >= self._next_dino_log:
|
||||
LOG.info(
|
||||
"dino_reid: embedded %d crops in %.1f ms (alpha=%.2f, "
|
||||
"matched %d mesh<->vision pairs)",
|
||||
len(v_crops), dt_ms, alpha, min(n_m, n_v))
|
||||
self._next_dino_log = now + 5.0
|
||||
|
||||
@staticmethod
|
||||
def _project_pelvis(
|
||||
translation: np.ndarray,
|
||||
|
||||
@@ -22,6 +22,7 @@ from pathlib import Path
|
||||
from .action_head_pub import ActionHeadPublisher
|
||||
from .euro_filter import SkeletonFilter
|
||||
from .pose_bridge import PoseSoundBridge
|
||||
from .pose_filter import PoseFilterChain
|
||||
from .state import Kp3D, PoseKp, State
|
||||
from .tracker import IoUTracker
|
||||
|
||||
@@ -96,6 +97,8 @@ class MultiWorker:
|
||||
self._sound_bridge = PoseSoundBridge(throttle_hz=30.0)
|
||||
self._action_pub = ActionHeadPublisher(state=self.state, bridge=self._sound_bridge)
|
||||
self._action_pub.start()
|
||||
# 3D pose filter chain : median, Kalman CV, lookahead, IK clamps.
|
||||
self._filter_chain = PoseFilterChain(state=self.state)
|
||||
|
||||
def start(self) -> None:
|
||||
self._thread = threading.Thread(
|
||||
@@ -251,6 +254,13 @@ class MultiWorker:
|
||||
# 3D world landmarks share ids with bodies (same MediaPipe
|
||||
# detection, just a different coordinate space).
|
||||
ids_body3d = ids_body[:len(bodies3d)] if bodies3d else []
|
||||
if bodies3d:
|
||||
bodies3d = self._filter_chain.apply(bodies3d, ids_body3d, t_now)
|
||||
# Debug : log body3d count once / 5 s so we know MediaPipe
|
||||
# actually populates pose_world_landmarks.
|
||||
if not hasattr(self, "_dbg_b3d_t") or t_now - self._dbg_b3d_t > 5.0:
|
||||
LOG.info("body3d: n=%d (pose_world_landmarks)", len(bodies3d))
|
||||
self._dbg_b3d_t = t_now
|
||||
self._sound_bridge.send(
|
||||
bodies, ids_body, t_now,
|
||||
persons_face=faces, persons_face_ids=ids_face,
|
||||
|
||||
@@ -238,6 +238,10 @@ class MultiHMRWorker:
|
||||
prev_thumb = thumb
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
# Publish to state for DINOv2 reid in MeshRigger.
|
||||
with self.state.lock():
|
||||
self.state.last_frame_rgb = frame_rgb
|
||||
self.state.last_frame_rgb_t = time.monotonic()
|
||||
tensor = torch.from_numpy(frame_rgb).permute(2, 0, 1).float()
|
||||
tensor = (tensor / 255.0).unsqueeze(0).to(device)
|
||||
|
||||
@@ -517,6 +521,9 @@ class MultiHMRWorker:
|
||||
prev_thumb = thumb
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
with self.state.lock():
|
||||
self.state.last_frame_rgb = frame_rgb
|
||||
self.state.last_frame_rgb_t = time.monotonic()
|
||||
img = frame_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
|
||||
|
||||
t_inf_start = time.monotonic()
|
||||
|
||||
@@ -0,0 +1,522 @@
|
||||
"""3D pose filtering chain : median spike removal, Kalman CV smoothing,
|
||||
spring-damper organic inertia, lookahead extrapolation, IK angular clamps.
|
||||
|
||||
Operates on lists of Kp3D (metric, hip-centered) keyed by track id.
|
||||
|
||||
Stages are toggleable via the POSE_FILTER env var :
|
||||
POSE_FILTER=median+kalman+lookahead+ik (default)
|
||||
POSE_FILTER=median
|
||||
POSE_FILTER=off
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Iterable
|
||||
|
||||
from .state import Kp3D, State
|
||||
|
||||
LOG = logging.getLogger("pose_filter")
|
||||
|
||||
NUM_JOINTS = 33
|
||||
DEFAULT_STAGES = ("median", "kalman", "lookahead", "ik")
|
||||
ALL_STAGES = ("median", "kalman", "spring", "lookahead", "ik")
|
||||
|
||||
# MediaPipe POSE_LANDMARKS indices used by IK constraints.
|
||||
L_SHOULDER, R_SHOULDER = 11, 12
|
||||
L_ELBOW, R_ELBOW = 13, 14
|
||||
L_WRIST, R_WRIST = 15, 16
|
||||
L_HIP, R_HIP = 23, 24
|
||||
L_KNEE, R_KNEE = 25, 26
|
||||
L_ANKLE, R_ANKLE = 27, 28
|
||||
L_FOOT, R_FOOT = 31, 32
|
||||
|
||||
# (parent_idx, joint_idx, child_idx, min_deg, max_deg)
|
||||
JOINT_LIMITS: tuple[tuple[int, int, int, float, float], ...] = (
|
||||
(L_SHOULDER, L_ELBOW, L_WRIST, 0.0, 175.0),
|
||||
(R_SHOULDER, R_ELBOW, R_WRIST, 0.0, 175.0),
|
||||
(L_HIP, L_KNEE, L_ANKLE, 0.0, 175.0),
|
||||
(R_HIP, R_KNEE, R_ANKLE, 0.0, 175.0),
|
||||
(L_KNEE, L_ANKLE, L_FOOT, 60.0, 135.0),
|
||||
(R_KNEE, R_ANKLE, R_FOOT, 60.0, 135.0),
|
||||
)
|
||||
|
||||
|
||||
# ----------------------------- utilities --------------------------------
|
||||
|
||||
def _is_finite(v: float) -> bool:
|
||||
return v == v and v not in (float("inf"), float("-inf"))
|
||||
|
||||
|
||||
def _kp_finite(kp: Kp3D) -> bool:
|
||||
return _is_finite(kp.x) and _is_finite(kp.y) and _is_finite(kp.z)
|
||||
|
||||
|
||||
def _median(values: list[float]) -> float:
|
||||
s = sorted(values)
|
||||
n = len(s)
|
||||
if n == 0:
|
||||
return 0.0
|
||||
if n % 2 == 1:
|
||||
return s[n // 2]
|
||||
return 0.5 * (s[n // 2 - 1] + s[n // 2])
|
||||
|
||||
|
||||
def _std(values: list[float], mu: float) -> float:
|
||||
if not values:
|
||||
return 0.0
|
||||
var = sum((v - mu) ** 2 for v in values) / len(values)
|
||||
return math.sqrt(var)
|
||||
|
||||
|
||||
# ----------------------------- median filter ----------------------------
|
||||
|
||||
class MedianFilter:
|
||||
"""Per (pid, joint) ring buffer ; replaces spikes outside 3σ by median."""
|
||||
|
||||
def __init__(self, window: int = 3) -> None:
|
||||
self.window = max(1, window)
|
||||
self._buf: dict[tuple[int, int], deque[tuple[float, float, float]]] = {}
|
||||
|
||||
def reset(self) -> None:
|
||||
self._buf.clear()
|
||||
|
||||
def apply(self, pid: int, joint_idx: int, x: float, y: float, z: float
|
||||
) -> tuple[float, float, float]:
|
||||
key = (pid, joint_idx)
|
||||
buf = self._buf.get(key)
|
||||
if buf is None:
|
||||
buf = deque(maxlen=self.window)
|
||||
self._buf[key] = buf
|
||||
|
||||
# Spike detection requires history.
|
||||
out = (x, y, z)
|
||||
if not (_is_finite(x) and _is_finite(y) and _is_finite(z)):
|
||||
if buf:
|
||||
med = (_median([v[0] for v in buf]),
|
||||
_median([v[1] for v in buf]),
|
||||
_median([v[2] for v in buf]))
|
||||
out = med
|
||||
else:
|
||||
out = (0.0, 0.0, 0.0)
|
||||
elif len(buf) >= self.window:
|
||||
for axis_idx, val in enumerate(out):
|
||||
col = [v[axis_idx] for v in buf]
|
||||
med = _median(col)
|
||||
sigma = _std(col, med)
|
||||
if sigma > 1e-6 and abs(val - med) > 3.0 * sigma:
|
||||
out = tuple(med if i == axis_idx else out[i]
|
||||
for i in range(3)) # type: ignore[assignment]
|
||||
buf.append(out)
|
||||
return out
|
||||
|
||||
|
||||
# ----------------------------- Kalman CV --------------------------------
|
||||
|
||||
@dataclass
|
||||
class _KalmanState:
|
||||
# State vector [x, y, z, vx, vy, vz]
|
||||
x: list[float] = field(default_factory=lambda: [0.0] * 6)
|
||||
# 6x6 covariance flattened
|
||||
P: list[list[float]] = field(default_factory=lambda: [[0.0] * 6 for _ in range(6)])
|
||||
initialised: bool = False
|
||||
last_t: float = 0.0
|
||||
|
||||
|
||||
def _mat_eye(n: int, s: float = 1.0) -> list[list[float]]:
|
||||
return [[s if i == j else 0.0 for j in range(n)] for i in range(n)]
|
||||
|
||||
|
||||
def _mat_mul(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
|
||||
ra, ca = len(A), len(A[0])
|
||||
cb = len(B[0])
|
||||
out = [[0.0] * cb for _ in range(ra)]
|
||||
for i in range(ra):
|
||||
Ai = A[i]
|
||||
for k in range(ca):
|
||||
aik = Ai[k]
|
||||
if aik == 0.0:
|
||||
continue
|
||||
Bk = B[k]
|
||||
for j in range(cb):
|
||||
out[i][j] += aik * Bk[j]
|
||||
return out
|
||||
|
||||
|
||||
def _mat_add(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
|
||||
return [[A[i][j] + B[i][j] for j in range(len(A[0]))] for i in range(len(A))]
|
||||
|
||||
|
||||
def _mat_sub(A: list[list[float]], B: list[list[float]]) -> list[list[float]]:
|
||||
return [[A[i][j] - B[i][j] for j in range(len(A[0]))] for i in range(len(A))]
|
||||
|
||||
|
||||
def _mat_T(A: list[list[float]]) -> list[list[float]]:
|
||||
return [[A[i][j] for i in range(len(A))] for j in range(len(A[0]))]
|
||||
|
||||
|
||||
def _mat_inv3(M: list[list[float]]) -> list[list[float]]:
|
||||
a, b, c = M[0]
|
||||
d, e, f = M[1]
|
||||
g, h, i = M[2]
|
||||
A = e * i - f * h
|
||||
B = -(d * i - f * g)
|
||||
C = d * h - e * g
|
||||
det = a * A + b * B + c * C
|
||||
if abs(det) < 1e-12:
|
||||
return _mat_eye(3, 1.0)
|
||||
inv_det = 1.0 / det
|
||||
return [
|
||||
[A * inv_det, -(b * i - c * h) * inv_det, (b * f - c * e) * inv_det],
|
||||
[B * inv_det, (a * i - c * g) * inv_det, -(a * f - c * d) * inv_det],
|
||||
[C * inv_det, -(a * h - b * g) * inv_det, (a * e - b * d) * inv_det],
|
||||
]
|
||||
|
||||
|
||||
class KalmanCV:
|
||||
"""Constant-velocity Kalman per (pid, joint_idx) on R^3."""
|
||||
|
||||
def __init__(self, q: float = 1e-3, r: float = 1e-2) -> None:
|
||||
self.q = q
|
||||
self.r = r
|
||||
self._states: dict[tuple[int, int], _KalmanState] = {}
|
||||
self._H = [
|
||||
[1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
|
||||
]
|
||||
|
||||
def reset(self) -> None:
|
||||
self._states.clear()
|
||||
|
||||
def get_velocity(self, pid: int, joint_idx: int) -> tuple[float, float, float]:
|
||||
st = self._states.get((pid, joint_idx))
|
||||
if st is None or not st.initialised:
|
||||
return (0.0, 0.0, 0.0)
|
||||
return (st.x[3], st.x[4], st.x[5])
|
||||
|
||||
def step(self, pid: int, joint_idx: int, mx: float, my: float, mz: float,
|
||||
t_now: float) -> tuple[float, float, float]:
|
||||
key = (pid, joint_idx)
|
||||
st = self._states.get(key)
|
||||
if st is None:
|
||||
st = _KalmanState()
|
||||
self._states[key] = st
|
||||
|
||||
if not st.initialised:
|
||||
st.x = [mx, my, mz, 0.0, 0.0, 0.0]
|
||||
st.P = _mat_eye(6, 1.0)
|
||||
st.initialised = True
|
||||
st.last_t = t_now
|
||||
return (mx, my, mz)
|
||||
|
||||
dt = max(1e-3, min(0.2, t_now - st.last_t))
|
||||
st.last_t = t_now
|
||||
|
||||
# Predict
|
||||
F = _mat_eye(6, 1.0)
|
||||
F[0][3] = dt
|
||||
F[1][4] = dt
|
||||
F[2][5] = dt
|
||||
x_pred = [
|
||||
st.x[0] + dt * st.x[3],
|
||||
st.x[1] + dt * st.x[4],
|
||||
st.x[2] + dt * st.x[5],
|
||||
st.x[3], st.x[4], st.x[5],
|
||||
]
|
||||
Q = _mat_eye(6, self.q)
|
||||
P_pred = _mat_add(_mat_mul(_mat_mul(F, st.P), _mat_T(F)), Q)
|
||||
|
||||
# Update
|
||||
z = [mx, my, mz]
|
||||
# y = z - H x_pred
|
||||
Hx = [x_pred[0], x_pred[1], x_pred[2]]
|
||||
y = [z[i] - Hx[i] for i in range(3)]
|
||||
# S = H P H^T + R (3x3)
|
||||
HP = _mat_mul(self._H, P_pred)
|
||||
S = [[HP[i][j] for j in range(3)] for i in range(3)]
|
||||
# add HP*H^T rest cols (cols 3..5) -> 0 contribution since H rest zero
|
||||
for i in range(3):
|
||||
S[i][i] += self.r
|
||||
S_inv = _mat_inv3(S)
|
||||
# K = P H^T S^-1 (6x3)
|
||||
PHt = [[P_pred[i][j] for j in range(3)] for i in range(6)]
|
||||
K = _mat_mul(PHt, S_inv)
|
||||
# x = x_pred + K y
|
||||
x_new = [x_pred[i] + sum(K[i][j] * y[j] for j in range(3))
|
||||
for i in range(6)]
|
||||
# P = (I - K H) P_pred
|
||||
KH = [[K[i][0] if j == 0 else (K[i][1] if j == 1 else (K[i][2] if j == 2 else 0.0))
|
||||
for j in range(6)] for i in range(6)]
|
||||
I6 = _mat_eye(6, 1.0)
|
||||
st.P = _mat_mul(_mat_sub(I6, KH), P_pred)
|
||||
st.x = x_new
|
||||
return (x_new[0], x_new[1], x_new[2])
|
||||
|
||||
|
||||
# --------------------------- spring damper ------------------------------
|
||||
|
||||
class SpringDamper:
|
||||
"""Critically-tunable spring-damper per (pid, joint_idx) on R^3."""
|
||||
|
||||
def __init__(self, stiffness: float = 200.0, damping: float = 15.0,
|
||||
mass: float = 1.0, enabled: bool = True) -> None:
|
||||
self.k = stiffness
|
||||
self.c = damping
|
||||
self.m = max(1e-3, mass)
|
||||
self.enabled = enabled
|
||||
self._pos: dict[tuple[int, int], list[float]] = {}
|
||||
self._vel: dict[tuple[int, int], list[float]] = {}
|
||||
self._last_t: dict[tuple[int, int], float] = {}
|
||||
|
||||
def reset(self) -> None:
|
||||
self._pos.clear()
|
||||
self._vel.clear()
|
||||
self._last_t.clear()
|
||||
|
||||
def step(self, pid: int, joint_idx: int, tx: float, ty: float, tz: float,
|
||||
t_now: float) -> tuple[float, float, float]:
|
||||
if not self.enabled:
|
||||
return (tx, ty, tz)
|
||||
key = (pid, joint_idx)
|
||||
pos = self._pos.get(key)
|
||||
if pos is None:
|
||||
self._pos[key] = [tx, ty, tz]
|
||||
self._vel[key] = [0.0, 0.0, 0.0]
|
||||
self._last_t[key] = t_now
|
||||
return (tx, ty, tz)
|
||||
dt = max(1e-3, min(0.1, t_now - self._last_t[key]))
|
||||
self._last_t[key] = t_now
|
||||
vel = self._vel[key]
|
||||
target = (tx, ty, tz)
|
||||
for i in range(3):
|
||||
# F = k(target - pos) - c * vel
|
||||
f = self.k * (target[i] - pos[i]) - self.c * vel[i]
|
||||
a = f / self.m
|
||||
vel[i] += a * dt
|
||||
pos[i] += vel[i] * dt
|
||||
return (pos[0], pos[1], pos[2])
|
||||
|
||||
|
||||
# --------------------------- lookahead ----------------------------------
|
||||
|
||||
class LookaheadPredictor:
|
||||
"""Linear extrapolation using Kalman velocities, capped to avoid blow-ups."""
|
||||
|
||||
def __init__(self, lookahead_ms: float = 50.0, max_velocity: float = 5.0
|
||||
) -> None:
|
||||
self.lookahead_s = lookahead_ms / 1000.0
|
||||
self.max_v = max_velocity
|
||||
|
||||
def step(self, x: float, y: float, z: float,
|
||||
vx: float, vy: float, vz: float) -> tuple[float, float, float]:
|
||||
def clamp(v: float) -> float:
|
||||
if v > self.max_v:
|
||||
return self.max_v
|
||||
if v < -self.max_v:
|
||||
return -self.max_v
|
||||
return v
|
||||
dt = self.lookahead_s
|
||||
return (x + clamp(vx) * dt, y + clamp(vy) * dt, z + clamp(vz) * dt)
|
||||
|
||||
|
||||
# --------------------------- IK constraints -----------------------------
|
||||
|
||||
def _vec_sub(a: tuple[float, float, float], b: tuple[float, float, float]
|
||||
) -> tuple[float, float, float]:
|
||||
return (a[0] - b[0], a[1] - b[1], a[2] - b[2])
|
||||
|
||||
|
||||
def _vec_add(a: tuple[float, float, float], b: tuple[float, float, float]
|
||||
) -> tuple[float, float, float]:
|
||||
return (a[0] + b[0], a[1] + b[1], a[2] + b[2])
|
||||
|
||||
|
||||
def _vec_scale(a: tuple[float, float, float], s: float
|
||||
) -> tuple[float, float, float]:
|
||||
return (a[0] * s, a[1] * s, a[2] * s)
|
||||
|
||||
|
||||
def _vec_dot(a: tuple[float, float, float], b: tuple[float, float, float]
|
||||
) -> float:
|
||||
return a[0] * b[0] + a[1] * b[1] + a[2] * b[2]
|
||||
|
||||
|
||||
def _vec_norm(a: tuple[float, float, float]) -> float:
|
||||
return math.sqrt(_vec_dot(a, a))
|
||||
|
||||
|
||||
def _vec_normalize(a: tuple[float, float, float], eps: float = 1e-9
|
||||
) -> tuple[float, float, float]:
|
||||
n = _vec_norm(a)
|
||||
if n < eps:
|
||||
return (1.0, 0.0, 0.0)
|
||||
return (a[0] / n, a[1] / n, a[2] / n)
|
||||
|
||||
|
||||
def _slerp_dir(d_from: tuple[float, float, float],
|
||||
d_to: tuple[float, float, float],
|
||||
t: float) -> tuple[float, float, float]:
|
||||
"""Slerp between two unit-ish vectors."""
|
||||
a = _vec_normalize(d_from)
|
||||
b = _vec_normalize(d_to)
|
||||
cos_a = max(-1.0, min(1.0, _vec_dot(a, b)))
|
||||
ang = math.acos(cos_a)
|
||||
if ang < 1e-6:
|
||||
return a
|
||||
sa = math.sin(ang)
|
||||
if abs(sa) < 1e-6:
|
||||
# antiparallel : pick an arbitrary perpendicular, then rotate.
|
||||
ortho = (1.0, 0.0, 0.0) if abs(a[0]) < 0.9 else (0.0, 1.0, 0.0)
|
||||
# Gram-Schmidt
|
||||
d = _vec_dot(ortho, a)
|
||||
perp = (ortho[0] - d * a[0], ortho[1] - d * a[1], ortho[2] - d * a[2])
|
||||
perp = _vec_normalize(perp)
|
||||
# rotate a by t*pi around perp axis : Rodrigues for angle = t*pi
|
||||
theta = t * ang
|
||||
cs, sn = math.cos(theta), math.sin(theta)
|
||||
# cross(perp, a)
|
||||
cx = perp[1] * a[2] - perp[2] * a[1]
|
||||
cy = perp[2] * a[0] - perp[0] * a[2]
|
||||
cz = perp[0] * a[1] - perp[1] * a[0]
|
||||
dot_pa = _vec_dot(perp, a)
|
||||
return (a[0] * cs + cx * sn + perp[0] * dot_pa * (1 - cs),
|
||||
a[1] * cs + cy * sn + perp[1] * dot_pa * (1 - cs),
|
||||
a[2] * cs + cz * sn + perp[2] * dot_pa * (1 - cs))
|
||||
w1 = math.sin((1.0 - t) * ang) / sa
|
||||
w2 = math.sin(t * ang) / sa
|
||||
return (a[0] * w1 + b[0] * w2,
|
||||
a[1] * w1 + b[1] * w2,
|
||||
a[2] * w1 + b[2] * w2)
|
||||
|
||||
|
||||
class IKConstraints:
|
||||
"""Clamp interior joint angles for elbows, knees, ankles."""
|
||||
|
||||
def __init__(self, limits: Iterable[tuple[int, int, int, float, float]]
|
||||
= JOINT_LIMITS) -> None:
|
||||
self.limits = tuple(limits)
|
||||
|
||||
def apply(self, kps: list[Kp3D]) -> list[Kp3D]:
|
||||
if len(kps) < NUM_JOINTS:
|
||||
return kps
|
||||
out = list(kps)
|
||||
for parent_i, joint_i, child_i, min_deg, max_deg in self.limits:
|
||||
if max(parent_i, joint_i, child_i) >= len(out):
|
||||
continue
|
||||
p = (out[parent_i].x, out[parent_i].y, out[parent_i].z)
|
||||
j = (out[joint_i].x, out[joint_i].y, out[joint_i].z)
|
||||
c = (out[child_i].x, out[child_i].y, out[child_i].z)
|
||||
v_pj = _vec_sub(p, j) # from joint to parent
|
||||
v_cj = _vec_sub(c, j) # from joint to child
|
||||
n_pj = _vec_norm(v_pj)
|
||||
n_cj = _vec_norm(v_cj)
|
||||
if n_pj < 1e-6 or n_cj < 1e-6:
|
||||
continue
|
||||
cos_a = max(-1.0, min(1.0, _vec_dot(v_pj, v_cj) / (n_pj * n_cj)))
|
||||
ang_deg = math.degrees(math.acos(cos_a))
|
||||
min_r = math.radians(min_deg)
|
||||
max_r = math.radians(max_deg)
|
||||
target_r: float | None = None
|
||||
if ang_deg < min_deg:
|
||||
target_r = min_r
|
||||
elif ang_deg > max_deg:
|
||||
target_r = max_r
|
||||
if target_r is None:
|
||||
continue
|
||||
# Interpolate child direction toward parent direction (or away)
|
||||
# so the new angle matches target_r.
|
||||
cur_r = math.acos(cos_a)
|
||||
# t such that new_angle = (1-t)*cur + t*pi between dirs ; use slerp.
|
||||
# Find t in [0,1] s.t. slerp(d_cj, d_pj, t) makes angle = target_r
|
||||
# The angle between slerp result and d_pj is (1-t)*cur_r.
|
||||
# So target_r = (1 - t) * cur_r -> t = 1 - target_r / cur_r
|
||||
if cur_r < 1e-6:
|
||||
continue
|
||||
t = 1.0 - (target_r / cur_r)
|
||||
t = max(0.0, min(1.0, t))
|
||||
d_cj = _vec_normalize(v_cj)
|
||||
d_pj = _vec_normalize(v_pj)
|
||||
new_dir = _slerp_dir(d_cj, d_pj, t)
|
||||
new_child = _vec_add(j, _vec_scale(new_dir, n_cj))
|
||||
old = out[child_i]
|
||||
out[child_i] = Kp3D(x=new_child[0], y=new_child[1],
|
||||
z=new_child[2], c=old.c)
|
||||
return out
|
||||
|
||||
|
||||
# --------------------------- chain wrapper ------------------------------
|
||||
|
||||
def _parse_env_stages() -> tuple[str, ...]:
|
||||
raw = os.environ.get("POSE_FILTER")
|
||||
if raw is None:
|
||||
return DEFAULT_STAGES
|
||||
raw = raw.strip().lower()
|
||||
if raw in ("off", "none", "0", "false"):
|
||||
return ()
|
||||
parts = tuple(p.strip() for p in raw.replace(",", "+").split("+") if p.strip())
|
||||
return tuple(p for p in parts if p in ALL_STAGES)
|
||||
|
||||
|
||||
class PoseFilterChain:
|
||||
"""Chain : median → kalman → spring → lookahead → ik."""
|
||||
|
||||
def __init__(self, state: State | None = None,
|
||||
enabled_stages: Iterable[str] | None = None) -> None:
|
||||
self.state = state
|
||||
if enabled_stages is None:
|
||||
stages = _parse_env_stages()
|
||||
else:
|
||||
stages = tuple(s for s in enabled_stages if s in ALL_STAGES)
|
||||
self.enabled = stages
|
||||
self.median = MedianFilter(window=3)
|
||||
self.kalman = KalmanCV()
|
||||
self.spring = SpringDamper(enabled="spring" in self.enabled)
|
||||
self.lookahead = LookaheadPredictor()
|
||||
self.ik = IKConstraints()
|
||||
self.last_apply_ms: float = 0.0
|
||||
LOG.info("PoseFilterChain stages=%s", self.enabled or ("off",))
|
||||
|
||||
def reset(self) -> None:
|
||||
self.median.reset()
|
||||
self.kalman.reset()
|
||||
self.spring.reset()
|
||||
|
||||
def apply(self, bodies3d: list[list[Kp3D]], ids: list[int],
|
||||
t_now: float) -> list[list[Kp3D]]:
|
||||
if not bodies3d or not self.enabled:
|
||||
self.last_apply_ms = 0.0
|
||||
return bodies3d
|
||||
t0 = time.perf_counter()
|
||||
out: list[list[Kp3D]] = []
|
||||
use_median = "median" in self.enabled
|
||||
use_kalman = "kalman" in self.enabled
|
||||
use_spring = "spring" in self.enabled
|
||||
use_lookahead = "lookahead" in self.enabled
|
||||
use_ik = "ik" in self.enabled
|
||||
|
||||
for body_i, kps in enumerate(bodies3d):
|
||||
pid = ids[body_i] if body_i < len(ids) else -1
|
||||
new_kps: list[Kp3D] = []
|
||||
for j_idx, kp in enumerate(kps):
|
||||
x, y, z, c = kp.x, kp.y, kp.z, kp.c
|
||||
if use_median:
|
||||
x, y, z = self.median.apply(pid, j_idx, x, y, z)
|
||||
if use_kalman:
|
||||
x, y, z = self.kalman.step(pid, j_idx, x, y, z, t_now)
|
||||
if use_spring:
|
||||
x, y, z = self.spring.step(pid, j_idx, x, y, z, t_now)
|
||||
if use_lookahead and use_kalman:
|
||||
vx, vy, vz = self.kalman.get_velocity(pid, j_idx)
|
||||
x, y, z = self.lookahead.step(x, y, z, vx, vy, vz)
|
||||
new_kps.append(Kp3D(x=x, y=y, z=z, c=c))
|
||||
if use_ik:
|
||||
new_kps = self.ik.apply(new_kps)
|
||||
out.append(new_kps)
|
||||
|
||||
self.last_apply_ms = (time.perf_counter() - t0) * 1000.0
|
||||
return out
|
||||
@@ -0,0 +1,202 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Convert DINOv2 ViT-S/14 to a CoreML .mlpackage for ANE-friendly inference.
|
||||
|
||||
The wrapped module takes (1, 3, 224, 224) RGB float32 in [0, 1], applies
|
||||
ImageNet normalization internally, runs the ViT, and returns the CLS
|
||||
embedding (1, 384) L2-normalised. We trace + convert with
|
||||
``coremltools.convert(... compute_units=ComputeUnit.ALL, compute_precision=FP16)``.
|
||||
|
||||
Run with the Python 3.12 venv that has coremltools and torch::
|
||||
|
||||
/tmp/coreml312/bin/python -m data_only_viz.scripts.convert_dinov2 [--force]
|
||||
|
||||
Output:
|
||||
~/.cache/av-live-multihmr/dinov2_vits14.mlpackage
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
LOG = logging.getLogger("convert_dinov2")
|
||||
|
||||
OUT_DIR = Path.home() / ".cache" / "av-live-multihmr"
|
||||
OUT_PATH = OUT_DIR / "dinov2_vits14.mlpackage"
|
||||
|
||||
_IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
||||
_IMAGENET_STD = (0.229, 0.224, 0.225)
|
||||
|
||||
|
||||
def _build_wrapper():
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
backbone = torch.hub.load(
|
||||
"facebookresearch/dinov2",
|
||||
"dinov2_vits14",
|
||||
source="github",
|
||||
trust_repo=True,
|
||||
)
|
||||
backbone.eval()
|
||||
|
||||
# Pretrained pos_embed is at 37x37 (518/14). We pre-resample to
|
||||
# 16x16 (224/14) once so the traced graph never needs an upsample.
|
||||
pe = backbone.pos_embed.data # (1, 1+37*37, 384)
|
||||
cls_pe = pe[:, :1]
|
||||
patch_pe = pe[:, 1:]
|
||||
n_old = int(round((patch_pe.shape[1]) ** 0.5))
|
||||
dim = patch_pe.shape[-1]
|
||||
patch_pe = patch_pe.reshape(1, n_old, n_old, dim).permute(0, 3, 1, 2)
|
||||
patch_pe = F.interpolate(patch_pe, size=(16, 16), mode="bilinear",
|
||||
align_corners=False)
|
||||
patch_pe = patch_pe.permute(0, 2, 3, 1).reshape(1, 16 * 16, dim)
|
||||
new_pe = torch.cat([cls_pe, patch_pe], dim=1).contiguous()
|
||||
backbone.pos_embed = nn.Parameter(new_pe, requires_grad=False)
|
||||
|
||||
mean = torch.tensor(_IMAGENET_MEAN, dtype=torch.float32).view(1, 3, 1, 1)
|
||||
std = torch.tensor(_IMAGENET_STD, dtype=torch.float32).view(1, 3, 1, 1)
|
||||
|
||||
class DinoV2Wrapper(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.backbone = backbone
|
||||
self.register_buffer("mean", mean)
|
||||
self.register_buffer("std", std)
|
||||
|
||||
def forward(self, x):
|
||||
x = (x - self.mean) / self.std
|
||||
bb = self.backbone
|
||||
x = bb.patch_embed(x)
|
||||
# cls_token is (1,1,384). Concat directly (B=1 fixed).
|
||||
x = torch.cat((bb.cls_token, x), dim=1)
|
||||
x = x + bb.pos_embed
|
||||
for blk in bb.blocks:
|
||||
x = blk(x)
|
||||
x = bb.norm(x)
|
||||
cls = x[:, 0]
|
||||
cls = cls / (cls.norm(dim=-1, keepdim=True) + 1e-8)
|
||||
return cls
|
||||
|
||||
return DinoV2Wrapper().eval()
|
||||
|
||||
|
||||
def _patch_coremltools_cast():
|
||||
"""coremltools 9.0 _cast assumes x.val is a 0-d scalar. With recent
|
||||
torch (2.12) some aten::Int args land as 1-D length-1 arrays. Patch
|
||||
the helper to flatten before scalar-casting."""
|
||||
from coremltools.converters.mil.frontend.torch import ops as _ops
|
||||
from coremltools.converters.mil.mil import Builder as mb
|
||||
|
||||
_orig = _ops._cast
|
||||
|
||||
def _patched_cast(context, node, dtype, dtype_name):
|
||||
# Inputs are read inside _orig from context; we wrap the failure
|
||||
# path by checking the first input's val first.
|
||||
inputs = _ops._get_inputs(context, node, expected=1)
|
||||
x = inputs[0]
|
||||
if x.can_be_folded_to_const():
|
||||
val = x.val
|
||||
if hasattr(val, "shape") and getattr(val, "shape", ()) != ():
|
||||
# 1-D length-1 (or all-ones shape) -> extract scalar
|
||||
import numpy as _np
|
||||
arr = _np.asarray(val).reshape(-1)
|
||||
if arr.size == 1:
|
||||
res = mb.const(val=dtype(arr[0]), name=node.name)
|
||||
context.add(res, node.name)
|
||||
return
|
||||
return _orig(context, node, dtype, dtype_name)
|
||||
|
||||
_ops._cast = _patched_cast
|
||||
|
||||
|
||||
def convert(force: bool = False) -> Path:
|
||||
import torch
|
||||
import coremltools as ct
|
||||
_patch_coremltools_cast()
|
||||
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
if OUT_PATH.exists() and not force:
|
||||
LOG.info("already converted: %s", OUT_PATH)
|
||||
return OUT_PATH
|
||||
|
||||
LOG.info("loading DINOv2 ViT-S/14 ...")
|
||||
wrap = _build_wrapper()
|
||||
example = torch.rand(1, 3, 224, 224, dtype=torch.float32)
|
||||
with torch.no_grad():
|
||||
ref_out = wrap(example)
|
||||
LOG.info("torch out shape=%s norm=%.4f", tuple(ref_out.shape),
|
||||
float(ref_out.norm(dim=-1).mean()))
|
||||
|
||||
LOG.info("tracing ...")
|
||||
with torch.no_grad():
|
||||
traced = torch.jit.trace(wrap, example, strict=False)
|
||||
|
||||
LOG.info("ct.convert (mlprogram FP16, computeUnits=ALL) ...")
|
||||
mlmodel = ct.convert(
|
||||
traced,
|
||||
source="pytorch",
|
||||
convert_to="mlprogram",
|
||||
inputs=[ct.TensorType(name="image", shape=example.shape,
|
||||
dtype=np.float32)],
|
||||
outputs=[ct.TensorType(name="embedding", dtype=np.float32)],
|
||||
compute_precision=ct.precision.FLOAT16,
|
||||
compute_units=ct.ComputeUnit.ALL,
|
||||
minimum_deployment_target=ct.target.macOS14,
|
||||
)
|
||||
mlmodel.short_description = "DINOv2 ViT-S/14 person re-id (384-D, L2)"
|
||||
mlmodel.save(str(OUT_PATH))
|
||||
LOG.info("saved %s", OUT_PATH)
|
||||
|
||||
pred = mlmodel.predict({"image": example.numpy().astype(np.float32)})
|
||||
coreml_out = list(pred.values())[0].reshape(-1)
|
||||
ref_np = ref_out.numpy().reshape(-1)
|
||||
cos = float(np.dot(coreml_out, ref_np) /
|
||||
(np.linalg.norm(coreml_out) * np.linalg.norm(ref_np) + 1e-8))
|
||||
LOG.info("CoreML vs Torch cosine on random input: %.4f", cos)
|
||||
return OUT_PATH
|
||||
|
||||
|
||||
def bench(n_iter: int = 30) -> None:
|
||||
import coremltools as ct
|
||||
LOG.info("bench: load mlpackage ...")
|
||||
m = ct.models.MLModel(str(OUT_PATH),
|
||||
compute_units=ct.ComputeUnit.ALL)
|
||||
crop = np.random.rand(1, 3, 224, 224).astype(np.float32)
|
||||
for _ in range(3):
|
||||
m.predict({"image": crop})
|
||||
times = []
|
||||
for _ in range(n_iter):
|
||||
t0 = time.perf_counter()
|
||||
m.predict({"image": crop})
|
||||
times.append((time.perf_counter() - t0) * 1e3)
|
||||
times.sort()
|
||||
p50 = times[len(times) // 2]
|
||||
p95 = times[int(len(times) * 0.95)]
|
||||
LOG.info("bench %d iter: p50=%.2f ms p95=%.2f ms mean=%.2f ms (~%.1f fps)",
|
||||
n_iter, p50, p95, sum(times) / len(times), 1000.0 / p50)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format="%(asctime)s %(name)s %(message)s")
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--force", action="store_true")
|
||||
ap.add_argument("--bench-only", action="store_true")
|
||||
ap.add_argument("--n-iter", type=int, default=30)
|
||||
args = ap.parse_args()
|
||||
|
||||
if not args.bench_only:
|
||||
convert(force=args.force)
|
||||
bench(n_iter=args.n_iter)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -25,9 +25,16 @@ from typing import Sequence
|
||||
|
||||
import numpy as np
|
||||
|
||||
import os
|
||||
|
||||
from .mesh_rigger import MeshRigger
|
||||
from .state import SMPLXPerson, State
|
||||
|
||||
try:
|
||||
from .dino_reid import DinoReid
|
||||
except Exception: # noqa: BLE001
|
||||
DinoReid = None # type: ignore[assignment]
|
||||
|
||||
LOG = logging.getLogger("smplx_tcp")
|
||||
|
||||
MAGIC = b"SMPX"
|
||||
@@ -47,7 +54,25 @@ class SMPLXTCPSender:
|
||||
self._sock: socket.socket | None = None
|
||||
# Hybrid keyframe rigging : entre deux keyframes Multi-HMR (~3 fps),
|
||||
# on translate le mesh via le delta pelvis Apple Vision (30 fps).
|
||||
self._rigger = MeshRigger(state) if enable_rigging else None
|
||||
# MULTIHMR_REID: 'dino' (try DINOv2 + IoU fusion, fallback IoU) /
|
||||
# 'iou' (pure IoU). Default: 'dino' if mlpackage exists.
|
||||
reid_mode = os.environ.get("MULTIHMR_REID", "dino").lower()
|
||||
dino = None
|
||||
if enable_rigging and reid_mode == "dino" and DinoReid is not None:
|
||||
try:
|
||||
if DinoReid.is_available():
|
||||
dino = DinoReid()
|
||||
LOG.info("MeshRigger: DINOv2 reid enabled")
|
||||
else:
|
||||
LOG.info(
|
||||
"MeshRigger: dino mlpackage absent, IoU only")
|
||||
except Exception as e: # noqa: BLE001
|
||||
LOG.warning("MeshRigger: dino load failed (%s), IoU only", e)
|
||||
dino = None
|
||||
dino_weight = float(os.environ.get("MULTIHMR_REID_ALPHA", "0.5"))
|
||||
self._rigger = MeshRigger(
|
||||
state, dino_weight=dino_weight,
|
||||
dino_reid=dino) if enable_rigging else None
|
||||
|
||||
def start(self) -> None:
|
||||
self._thread = threading.Thread(
|
||||
|
||||
@@ -140,6 +140,11 @@ class State:
|
||||
# Derniere frame webcam au format JPEG bytes (pour NSImageView overlay).
|
||||
# Le pose worker la met a jour ; le HUD timer lit et l'affiche.
|
||||
last_webcam_jpeg: bytes | None = None
|
||||
# Last full RGB frame fed to Multi-HMR (uint8 HxWx3, typ. 672x672).
|
||||
# Updated by multi_hmr_worker right before inference. Read by
|
||||
# MeshRigger for DINOv2-based person re-id. None when absent.
|
||||
last_frame_rgb: np.ndarray | None = None
|
||||
last_frame_rgb_t: float = 0.0
|
||||
|
||||
_lock: threading.RLock = field(default_factory=threading.RLock, repr=False)
|
||||
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Tests for the DINOv2 reid backend.
|
||||
|
||||
These tests are skipped automatically if the .mlpackage is not present
|
||||
(`scripts/convert_dinov2.py` was never run) or pyobjc is unavailable.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from data_only_viz.dino_reid import DEFAULT_MLPACKAGE, EMBED_DIM, DinoReid
|
||||
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not DEFAULT_MLPACKAGE.exists(),
|
||||
reason=f"DINOv2 mlpackage missing at {DEFAULT_MLPACKAGE}; "
|
||||
"run scripts/convert_dinov2.py first",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def reid() -> DinoReid:
|
||||
return DinoReid()
|
||||
|
||||
|
||||
def test_is_available() -> None:
|
||||
assert DinoReid.is_available() is True
|
||||
|
||||
|
||||
def test_load(reid: DinoReid) -> None:
|
||||
assert reid is not None
|
||||
assert reid._out_name
|
||||
|
||||
|
||||
def test_embed_random_crops_different(reid: DinoReid) -> None:
|
||||
# Two crops with very different visual content. DINOv2 CLS tokens
|
||||
# for two iid noise patches are surprisingly close (~0.98), so we
|
||||
# build crops that are visually distinct: one is mostly red, the
|
||||
# other is mostly green with a striped pattern.
|
||||
a = np.zeros((224, 224, 3), dtype=np.uint8)
|
||||
a[..., 0] = 220 # red
|
||||
a[40:80, 40:180] = (240, 30, 30)
|
||||
b = np.zeros((224, 224, 3), dtype=np.uint8)
|
||||
b[..., 1] = 200 # green
|
||||
for i in range(0, 224, 16):
|
||||
b[i:i + 8] = (10, 30, 220) # blue stripes
|
||||
embs = reid.embed_crops([a, b])
|
||||
assert embs.shape == (2, EMBED_DIM)
|
||||
norms = np.linalg.norm(embs, axis=1)
|
||||
assert np.allclose(norms, 1.0, atol=1e-3)
|
||||
cos = float(np.dot(embs[0], embs[1]))
|
||||
assert cos < 0.95, f"distinct crops too similar: cos={cos:.3f}"
|
||||
|
||||
|
||||
def test_embed_identical_crops_same(reid: DinoReid) -> None:
|
||||
rng = np.random.default_rng(7)
|
||||
a = rng.integers(0, 255, size=(224, 224, 3), dtype=np.uint8)
|
||||
embs = reid.embed_crops([a, a.copy()])
|
||||
assert embs.shape == (2, EMBED_DIM)
|
||||
cos = float(np.dot(embs[0], embs[1]))
|
||||
assert cos > 0.999, f"identical crops cos={cos:.4f} (expected ~1.0)"
|
||||
|
||||
|
||||
def test_latency_batch4(reid: DinoReid) -> None:
|
||||
rng = np.random.default_rng(0)
|
||||
crops = [rng.integers(0, 255, size=(180, 90, 3), dtype=np.uint8)
|
||||
for _ in range(4)]
|
||||
# warmup
|
||||
reid.embed_crops(crops)
|
||||
t0 = time.perf_counter()
|
||||
reid.embed_crops(crops)
|
||||
dt_ms = (time.perf_counter() - t0) * 1e3
|
||||
# Spec target: < 30 ms for batch=4 on M5.
|
||||
assert dt_ms < 80.0, f"batch=4 too slow: {dt_ms:.1f} ms"
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Tests for the 3D pose filter chain."""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from data_only_viz.pose_filter import (
|
||||
IKConstraints,
|
||||
KalmanCV,
|
||||
LookaheadPredictor,
|
||||
MedianFilter,
|
||||
PoseFilterChain,
|
||||
L_ELBOW,
|
||||
L_SHOULDER,
|
||||
L_WRIST,
|
||||
)
|
||||
from data_only_viz.state import Kp3D
|
||||
|
||||
|
||||
def _body(values: list[tuple[float, float, float]]) -> list[Kp3D]:
|
||||
"""Build a 33-joint body, fill remaining with zeros."""
|
||||
out = [Kp3D(x=v[0], y=v[1], z=v[2], c=1.0) for v in values]
|
||||
while len(out) < 33:
|
||||
out.append(Kp3D(x=0.0, y=0.0, z=0.0, c=1.0))
|
||||
return out
|
||||
|
||||
|
||||
def test_median_filter_kills_spike() -> None:
|
||||
mf = MedianFilter(window=3)
|
||||
pid, j = 0, 0
|
||||
# Warm up
|
||||
mf.apply(pid, j, 0.0, 0.0, 0.0)
|
||||
mf.apply(pid, j, 0.01, 0.0, 0.0)
|
||||
mf.apply(pid, j, 0.02, 0.0, 0.0)
|
||||
# Spike (NaN)
|
||||
x, y, z = mf.apply(pid, j, float("nan"), float("nan"), float("nan"))
|
||||
assert math.isfinite(x) and math.isfinite(y) and math.isfinite(z)
|
||||
assert abs(x) < 0.1
|
||||
# Big outlier in x
|
||||
x2, _, _ = mf.apply(pid, j, 10.0, 0.0, 0.0)
|
||||
assert x2 < 1.0
|
||||
|
||||
|
||||
def test_kalman_converges() -> None:
|
||||
# Use a noisy constant-velocity signal : Kalman CV should converge.
|
||||
import random
|
||||
rng = random.Random(0)
|
||||
kf = KalmanCV(q=1e-3, r=1e-2)
|
||||
pid, j = 0, 0
|
||||
t = 0.0
|
||||
dt = 1.0 / 30.0
|
||||
vel = 0.3 # m/s
|
||||
errs: list[float] = []
|
||||
for i in range(120):
|
||||
t += dt
|
||||
true_pos = vel * t
|
||||
meas = true_pos + rng.gauss(0.0, 0.01) # 1 cm gaussian noise
|
||||
out = kf.step(pid, j, meas, 0.0, 0.0, t)
|
||||
if i > 30:
|
||||
errs.append(abs(out[0] - true_pos))
|
||||
mean_err = sum(errs) / len(errs)
|
||||
assert mean_err < 0.01 # ±1 cm post warmup
|
||||
|
||||
|
||||
def test_lookahead_extrapolates_constant_velocity() -> None:
|
||||
pred = LookaheadPredictor(lookahead_ms=50.0, max_velocity=5.0)
|
||||
x, y, z = pred.step(0.0, 0.0, 0.0, 1.0, 0.0, 0.0)
|
||||
assert abs(x - 0.05) < 1e-6
|
||||
assert abs(y) < 1e-9 and abs(z) < 1e-9
|
||||
# Velocity cap
|
||||
x2, _, _ = pred.step(0.0, 0.0, 0.0, 100.0, 0.0, 0.0)
|
||||
assert abs(x2 - 5.0 * 0.050) < 1e-6
|
||||
|
||||
|
||||
def test_ik_clamps_elbow_180_plus() -> None:
|
||||
ik = IKConstraints()
|
||||
# Shoulder at origin, elbow at (1,0,0), wrist BEHIND elbow at (2,0,0)
|
||||
# -> shoulder-elbow-wrist angle is 180 deg, exceeds 175 deg limit.
|
||||
coords: list[tuple[float, float, float]] = [(0.0, 0.0, 0.0)] * 33
|
||||
coords[L_SHOULDER] = (0.0, 0.0, 0.0)
|
||||
coords[L_ELBOW] = (1.0, 0.0, 0.0)
|
||||
coords[L_WRIST] = (2.0, 0.0, 0.0)
|
||||
body = _body(coords)
|
||||
out = ik.apply(body)
|
||||
p = (out[L_SHOULDER].x, out[L_SHOULDER].y, out[L_SHOULDER].z)
|
||||
e = (out[L_ELBOW].x, out[L_ELBOW].y, out[L_ELBOW].z)
|
||||
w = (out[L_WRIST].x, out[L_WRIST].y, out[L_WRIST].z)
|
||||
v_pj = (p[0] - e[0], p[1] - e[1], p[2] - e[2])
|
||||
v_cj = (w[0] - e[0], w[1] - e[1], w[2] - e[2])
|
||||
n_pj = math.sqrt(sum(c * c for c in v_pj))
|
||||
n_cj = math.sqrt(sum(c * c for c in v_cj))
|
||||
cos_a = (v_pj[0] * v_cj[0] + v_pj[1] * v_cj[1] + v_pj[2] * v_cj[2]
|
||||
) / (n_pj * n_cj)
|
||||
cos_a = max(-1.0, min(1.0, cos_a))
|
||||
ang_deg = math.degrees(math.acos(cos_a))
|
||||
assert ang_deg <= 175.5
|
||||
# Bone length preserved
|
||||
assert abs(n_cj - 1.0) < 1e-6
|
||||
|
||||
|
||||
def test_chain_no_op_when_disabled() -> None:
|
||||
chain = PoseFilterChain(enabled_stages=())
|
||||
body = _body([(0.1, 0.2, 0.3), (0.4, 0.5, 0.6)])
|
||||
out = chain.apply([body], [0], t_now=0.0)
|
||||
assert len(out) == 1
|
||||
for i in range(len(body)):
|
||||
assert out[0][i].x == body[i].x
|
||||
assert out[0][i].y == body[i].y
|
||||
assert out[0][i].z == body[i].z
|
||||
|
||||
|
||||
def test_chain_latency_under_2ms() -> None:
|
||||
chain = PoseFilterChain(
|
||||
enabled_stages=("median", "kalman", "lookahead", "ik"))
|
||||
body = _body([(i * 0.01, i * 0.02, i * 0.03) for i in range(33)])
|
||||
# Warm up internal state
|
||||
for k in range(5):
|
||||
chain.apply([body, body], [0, 1], t_now=k * 0.033)
|
||||
# Measure
|
||||
times: list[float] = []
|
||||
for k in range(30):
|
||||
chain.apply([body, body], [0, 1], t_now=(k + 5) * 0.033)
|
||||
times.append(chain.last_apply_ms)
|
||||
avg = sum(times) / len(times)
|
||||
# Generous bound for CI ; live target is <2 ms but allow 10 ms in tests.
|
||||
assert avg < 10.0
|
||||
@@ -101,6 +101,15 @@ struct BodyView: NSViewRepresentable {
|
||||
context.coordinator.sceneRenderer = scene
|
||||
context.coordinator.mtkView = mtkView
|
||||
context.coordinator.skeletonOverlay = SkeletonOverlay(parent: bodyAnchor)
|
||||
// Skeleton 3D RealityKit armature (33 spheres + 32 cylinders bones)
|
||||
// driven by /pose3d/* OSC from MediaPipe pose_world_landmarks.
|
||||
// Visible quand toggle showSkeleton ou vizMode==9 (openpos).
|
||||
let skel3dAnchor = AnchorEntity(world: SIMD3<Float>(0, 0, -2.5))
|
||||
arView.scene.addAnchor(skel3dAnchor)
|
||||
let skel3d = Skeleton3DRenderer()
|
||||
skel3d.attach(to: skel3dAnchor, listener: poseListener)
|
||||
context.coordinator.skel3dAnchor = skel3dAnchor
|
||||
context.coordinator.skel3d = skel3d
|
||||
context.coordinator.keyLight = key
|
||||
context.coordinator.fillLight = fill
|
||||
context.coordinator.rimLight = rim
|
||||
@@ -162,6 +171,9 @@ struct BodyView: NSViewRepresentable {
|
||||
let skelVisible = settings.vizMode == 9 || settings.showSkeleton
|
||||
c.skeletonOverlay?.update(persons: poseListener.persons,
|
||||
visible: skelVisible)
|
||||
// 3D RealityKit armature : show/hide root anchor in sync with
|
||||
// the same skelVisible signal as the 2D overlay.
|
||||
c.skel3dAnchor?.isEnabled = skelVisible
|
||||
// Pose -> scene uniforms : drive hands3d (mode 8) et openpos
|
||||
// (mode 9) avec la premiere personne detectee. Les wrists pilotent
|
||||
// hand_l/r ; pose_count alimente bg_fragment.
|
||||
@@ -211,6 +223,8 @@ struct BodyView: NSViewRepresentable {
|
||||
var sceneRenderer: SceneRenderer?
|
||||
var mtkView: MTKView?
|
||||
var skeletonOverlay: SkeletonOverlay?
|
||||
var skel3dAnchor: AnchorEntity?
|
||||
var skel3d: Skeleton3DRenderer?
|
||||
var kbMonitor: Any?
|
||||
|
||||
deinit {
|
||||
|
||||
@@ -26,8 +26,39 @@ final class PoseOSCListener: ObservableObject {
|
||||
var seenAt: TimeInterval = 0
|
||||
}
|
||||
|
||||
/// MediaPipe pose_world_landmarks : 33 keypoints in meters, hip-rel.
|
||||
/// MediaPipe convention : x=right, y=down, z=forward (away from cam).
|
||||
struct Pose3DFrame: Equatable {
|
||||
var pid: Int = -1
|
||||
var kps: [SIMD4<Float>] = Array(repeating: .zero, count: 33)
|
||||
var hasPoint: [Bool] = Array(repeating: false, count: 33)
|
||||
var seenAt: TimeInterval = 0
|
||||
}
|
||||
|
||||
/// 68 dlib-style facial landmarks (x,y normalises 0..1).
|
||||
struct FaceFrame: Equatable {
|
||||
var points: [SIMD2<Float>] = Array(repeating: .zero, count: 68)
|
||||
var hasPoint: [Bool] = Array(repeating: false, count: 68)
|
||||
var seenAt: TimeInterval = 0
|
||||
}
|
||||
|
||||
/// 21 MediaPipe hand landmarks per detected hand.
|
||||
struct HandFrame: Equatable {
|
||||
var side: Int = 0
|
||||
var points: [SIMD2<Float>] = Array(repeating: .zero, count: 21)
|
||||
var hasPoint: [Bool] = Array(repeating: false, count: 21)
|
||||
var seenAt: TimeInterval = 0
|
||||
}
|
||||
|
||||
@Published var persons: [Int: PoseFrame] = [:]
|
||||
@Published var count: Int = 0
|
||||
@Published var body3d: [Int: Pose3DFrame] = [:]
|
||||
@Published var body3dCount: Int = 0
|
||||
@Published var faces: [Int: FaceFrame] = [:]
|
||||
@Published var hands: [Int: HandFrame] = [:]
|
||||
@Published var faceCount: Int = 0
|
||||
@Published var handCountLeft: Int = 0
|
||||
@Published var handCountRight: Int = 0
|
||||
|
||||
private var listener: NWListener?
|
||||
|
||||
@@ -148,13 +179,77 @@ final class PoseOSCListener: ObservableObject {
|
||||
p.skeleton = skel
|
||||
p.seenAt = CFAbsoluteTimeGetCurrent()
|
||||
persons[Int(pid)] = p
|
||||
case "/face/count":
|
||||
if let n = args.first as? Int32 { faceCount = Int(n) }
|
||||
if faceCount == 0 { faces.removeAll(keepingCapacity: true) }
|
||||
case "/face/kp":
|
||||
guard args.count >= 6,
|
||||
let pid = args[0] as? Int32,
|
||||
let slot = args[1] as? Int32,
|
||||
let x = args[2] as? Float,
|
||||
let y = args[3] as? Float else { return }
|
||||
let s = Int(slot)
|
||||
guard s >= 0 && s < 68 else { return }
|
||||
var f = faces[Int(pid)] ?? FaceFrame()
|
||||
f.points[s] = SIMD2(x, y)
|
||||
f.hasPoint[s] = true
|
||||
f.seenAt = CFAbsoluteTimeGetCurrent()
|
||||
faces[Int(pid)] = f
|
||||
case "/hand/count":
|
||||
if args.count >= 2,
|
||||
let l = args[0] as? Int32, let r = args[1] as? Int32 {
|
||||
handCountLeft = Int(l)
|
||||
handCountRight = Int(r)
|
||||
if handCountLeft + handCountRight == 0 {
|
||||
hands.removeAll(keepingCapacity: true)
|
||||
}
|
||||
}
|
||||
case "/hand/kp":
|
||||
guard args.count >= 7,
|
||||
let pid = args[0] as? Int32,
|
||||
let side = args[1] as? Int32,
|
||||
let idx = args[2] as? Int32,
|
||||
let x = args[3] as? Float,
|
||||
let y = args[4] as? Float else { return }
|
||||
let i = Int(idx)
|
||||
guard i >= 0 && i < 21 else { return }
|
||||
var h = hands[Int(pid)] ?? HandFrame()
|
||||
h.side = Int(side)
|
||||
h.points[i] = SIMD2(x, y)
|
||||
h.hasPoint[i] = true
|
||||
h.seenAt = CFAbsoluteTimeGetCurrent()
|
||||
hands[Int(pid)] = h
|
||||
case "/pose3d/count":
|
||||
if let n = args.first as? Int32 {
|
||||
body3dCount = Int(n)
|
||||
if body3dCount == 0 {
|
||||
body3d.removeAll(keepingCapacity: true)
|
||||
}
|
||||
}
|
||||
case "/pose3d/kp":
|
||||
guard args.count >= 6,
|
||||
let pid = args[0] as? Int32,
|
||||
let idx = args[1] as? Int32,
|
||||
let x = args[2] as? Float,
|
||||
let y = args[3] as? Float,
|
||||
let z = args[4] as? Float,
|
||||
let c = args[5] as? Float else { return }
|
||||
let i = Int(idx)
|
||||
guard i >= 0 && i < 33 else { return }
|
||||
var p = body3d[Int(pid)] ?? Pose3DFrame(pid: Int(pid))
|
||||
p.pid = Int(pid)
|
||||
p.kps[i] = SIMD4<Float>(x, y, z, c)
|
||||
p.hasPoint[i] = true
|
||||
p.seenAt = CFAbsoluteTimeGetCurrent()
|
||||
body3d[Int(pid)] = p
|
||||
default:
|
||||
break
|
||||
}
|
||||
// Garbage-collect persons non vues depuis > 2 s
|
||||
// Garbage-collect persons + body3d non vus depuis > 2 s
|
||||
let now = CFAbsoluteTimeGetCurrent()
|
||||
persons = persons.filter { $0.value.seenAt == 0
|
||||
|| now - $0.value.seenAt < 2.0 }
|
||||
body3d = body3d.filter { now - $0.value.seenAt < 2.0 }
|
||||
}
|
||||
|
||||
// MARK: - Minimal OSC parser
|
||||
|
||||
Reference in New Issue
Block a user