L'électron rare
7ed2e2764a
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.
205 lines
7.9 KiB
Python
205 lines
7.9 KiB
Python
"""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
|