feat(data-only-viz): Multi-HMR CoreML backend
Convert Multi-HMR ViT-S 672 to CoreML mlpackage and wire as optional worker backend via MULTIHMR_BACKEND=coreml env var. Inference path uses pyobjc + native CoreML framework (Python 3.14 has no libcoremlpython binding). Conversion done in a separate Py 3.12 venv; einsum cascade patched (camera intrinsics broadcast + smplx landmarks) via setup_multihmr.sh, idempotent on re-clone. Bench: 28 ms mock, 100-170 ms live (~13 fps, 4x PyTorch MPS). ANE compile fails on this model; CPU+GPU is the sweet spot.
This commit is contained in:
L'électron rare
@@ -30,6 +30,7 @@ CACHE = Path.home() / ".cache" / "av-live-multihmr"
|
||||
CKPT = CACHE / "checkpoints" / "multiHMR_672_S.pt"
|
||||
SMPLX_PATH = CACHE / "models" / "smplx" / "SMPLX_NEUTRAL.npz"
|
||||
MULTIHMR_REPO = CACHE / "multi-hmr"
|
||||
COREML_MLPACKAGE = CACHE / "multihmr_full_672_s.mlpackage"
|
||||
|
||||
IMG_SIZE = 672
|
||||
N_VERTS = 10475
|
||||
@@ -41,7 +42,8 @@ class MultiHMRWorker:
|
||||
det_thresh: float = 0.3,
|
||||
nms_kernel_size: int = 5,
|
||||
motion_gate: float = 5.0,
|
||||
camera_index: int = -1) -> None:
|
||||
camera_index: int = -1,
|
||||
backend: str | None = None) -> None:
|
||||
self.state = state
|
||||
self.num_persons = num_persons
|
||||
self.period = 1.0 / max(1.0, target_fps)
|
||||
@@ -55,6 +57,12 @@ class MultiHMRWorker:
|
||||
self.motion_gate = motion_gate
|
||||
# -1 = auto-select Mac BuiltInWideAngleCamera (cf _camera_select)
|
||||
self.camera_index = camera_index
|
||||
# backend: 'pytorch' (default) or 'coreml'. CoreML uses the
|
||||
# .mlpackage at COREML_MLPACKAGE, bypasses MPS torch, and runs
|
||||
# on ANE/GPU/CPU via CoreML.framework natively (3-4x faster).
|
||||
self.backend = (backend
|
||||
or os.environ.get("MULTIHMR_BACKEND", "pytorch")
|
||||
).strip().lower()
|
||||
self._stop = threading.Event()
|
||||
self._thread: threading.Thread | None = None
|
||||
self._smooth_shape = [
|
||||
@@ -72,6 +80,9 @@ class MultiHMRWorker:
|
||||
|
||||
@staticmethod
|
||||
def is_available() -> bool:
|
||||
backend = os.environ.get("MULTIHMR_BACKEND", "pytorch").strip().lower()
|
||||
if backend == "coreml":
|
||||
return COREML_MLPACKAGE.exists()
|
||||
return CKPT.exists() and SMPLX_PATH.exists() and MULTIHMR_REPO.exists()
|
||||
|
||||
def start(self) -> None:
|
||||
@@ -83,6 +94,12 @@ class MultiHMRWorker:
|
||||
self._stop.set()
|
||||
|
||||
def _run(self) -> None:
|
||||
if self.backend == "coreml":
|
||||
self._run_coreml()
|
||||
return
|
||||
self._run_pytorch()
|
||||
|
||||
def _run_pytorch(self) -> None:
|
||||
if str(MULTIHMR_REPO) not in sys.path:
|
||||
sys.path.insert(0, str(MULTIHMR_REPO))
|
||||
# Multi-HMR demo.py tire pyrender / pyvista (OpenGL offscreen) et
|
||||
@@ -373,3 +390,223 @@ class MultiHMRWorker:
|
||||
|
||||
cap.stop()
|
||||
LOG.info("multi_hmr worker stopped")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# CoreML backend
|
||||
# ------------------------------------------------------------------
|
||||
def _run_coreml(self) -> None:
|
||||
"""CoreML inference path (ANE+GPU+CPU via Apple's framework).
|
||||
|
||||
Mirrors _run_pytorch but loads the .mlpackage via pyobjc + the
|
||||
CoreML.framework, bypassing torch/MPS entirely. ~3-4x faster
|
||||
on M5 (28.8ms median vs ~100ms with MPS)."""
|
||||
try:
|
||||
import cv2
|
||||
except ImportError as e:
|
||||
LOG.error("opencv-python missing: %s", e)
|
||||
return
|
||||
try:
|
||||
from .multihmr_coreml import MultiHMRCoreMLBackend
|
||||
backend = MultiHMRCoreMLBackend(COREML_MLPACKAGE)
|
||||
except Exception as e: # noqa: BLE001
|
||||
LOG.error("CoreML backend init failed: %s", e)
|
||||
return
|
||||
|
||||
focal = float(IMG_SIZE)
|
||||
K_np = np.array([[focal, 0.0, IMG_SIZE / 2.0],
|
||||
[0.0, focal, IMG_SIZE / 2.0],
|
||||
[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
from ._av_capture import (
|
||||
AVCapture, find_builtin_device, enumerate_devices)
|
||||
if self.camera_index >= 0:
|
||||
devs = enumerate_devices()
|
||||
if self.camera_index >= len(devs):
|
||||
LOG.error("camera_index %d hors de %d devices",
|
||||
self.camera_index, len(devs))
|
||||
return
|
||||
info = devs[self.camera_index]
|
||||
else:
|
||||
info = find_builtin_device()
|
||||
if info is None:
|
||||
LOG.error("aucune BuiltInWideAngleCamera trouvee")
|
||||
return
|
||||
cap = AVCapture(info)
|
||||
if not cap.start():
|
||||
LOG.error("AVCapture start failed pour %s", info["name"])
|
||||
return
|
||||
LOG.info("camera ouverte %s (%s) [coreml backend]",
|
||||
info["name"], info["type"])
|
||||
|
||||
frame_count = 0
|
||||
persons_count = 0
|
||||
skipped_static = 0
|
||||
next_heartbeat = time.monotonic() + 5.0
|
||||
prev_thumb: np.ndarray | None = None
|
||||
|
||||
while not self._stop.is_set():
|
||||
t_cap_start = time.monotonic()
|
||||
ok, frame_bgr = cap.read(timeout_s=0.5)
|
||||
if not ok or frame_bgr is None:
|
||||
time.sleep(self.period)
|
||||
continue
|
||||
|
||||
t_pre_start = time.monotonic()
|
||||
h, w = frame_bgr.shape[:2]
|
||||
if (h, w) != (IMG_SIZE, IMG_SIZE):
|
||||
side = min(h, w)
|
||||
y0 = (h - side) // 2
|
||||
x0 = (w - side) // 2
|
||||
frame_bgr = frame_bgr[y0:y0 + side, x0:x0 + side]
|
||||
frame_bgr = cv2.resize(frame_bgr, (IMG_SIZE, IMG_SIZE))
|
||||
|
||||
if self.motion_gate > 0:
|
||||
thumb = cv2.cvtColor(
|
||||
cv2.resize(frame_bgr, (112, 112)),
|
||||
cv2.COLOR_BGR2GRAY)
|
||||
if prev_thumb is not None:
|
||||
diff_mean = float(np.mean(
|
||||
cv2.absdiff(thumb, prev_thumb)))
|
||||
if diff_mean < self.motion_gate:
|
||||
prev_thumb = thumb
|
||||
skipped_static += 1
|
||||
time.sleep(self.period)
|
||||
continue
|
||||
prev_thumb = thumb
|
||||
|
||||
frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
|
||||
img = frame_rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
|
||||
|
||||
t_inf_start = time.monotonic()
|
||||
try:
|
||||
humans = backend.infer(img, K_np, det_thresh=self.det_thresh)
|
||||
except Exception as e: # noqa: BLE001
|
||||
LOG.warning("coreml inference failed: %s", e)
|
||||
time.sleep(self.period)
|
||||
continue
|
||||
|
||||
t_post_start = time.monotonic()
|
||||
t_now = time.monotonic()
|
||||
frame_count += 1
|
||||
persons_count += len(humans) if humans else 0
|
||||
if t_now >= next_heartbeat:
|
||||
fps = frame_count / 5.0
|
||||
avg = persons_count / max(1, frame_count)
|
||||
LOG.info(
|
||||
"hb[coreml]: %.1f fps, %.2f persons/frame, %d skipped",
|
||||
fps, avg, skipped_static)
|
||||
frame_count = 0
|
||||
persons_count = 0
|
||||
skipped_static = 0
|
||||
next_heartbeat = t_now + 5.0
|
||||
|
||||
if not humans:
|
||||
with self.state.lock():
|
||||
self.state.persons_smplx = []
|
||||
time.sleep(self.period)
|
||||
continue
|
||||
|
||||
# Dedup intra-frame (same logic as pytorch path).
|
||||
cand: list[tuple[
|
||||
float, float, float, float, float,
|
||||
np.ndarray, int]] = []
|
||||
for i, hh in enumerate(humans):
|
||||
v = hh["v3d"].detach().cpu().numpy()
|
||||
xmin = float(v[:, 0].min()); ymin = float(v[:, 1].min())
|
||||
xmax = float(v[:, 0].max()); ymax = float(v[:, 1].max())
|
||||
score = float(hh["scores"].item())
|
||||
pelv = hh["transl_pelvis"].detach().cpu().numpy(
|
||||
).flatten()[:3]
|
||||
cand.append((score, xmin, ymin, xmax, ymax, pelv, i))
|
||||
cand.sort(key=lambda c: -c[0])
|
||||
keep_idx: list[int] = []
|
||||
kept: list[tuple[float, float, float, float, np.ndarray]] = []
|
||||
for sc, x0, y0, x1, y1, pelv, src_i in cand:
|
||||
a_area = max(0.0, x1 - x0) * max(0.0, y1 - y0)
|
||||
drop = False
|
||||
for (kx0, ky0, kx1, ky1, kpelv) in kept:
|
||||
ix0 = max(x0, kx0); iy0 = max(y0, ky0)
|
||||
ix1 = min(x1, kx1); iy1 = min(y1, ky1)
|
||||
iw = max(0.0, ix1 - ix0); ih = max(0.0, iy1 - iy0)
|
||||
inter = iw * ih
|
||||
if a_area <= 0 or inter <= 0:
|
||||
continue
|
||||
k_area = (kx1 - kx0) * (ky1 - ky0)
|
||||
iou = inter / (a_area + k_area - inter + 1e-9)
|
||||
pelv_d = float(np.linalg.norm(pelv - kpelv))
|
||||
if iou > 0.55 and pelv_d < 0.20:
|
||||
drop = True
|
||||
break
|
||||
if not drop:
|
||||
keep_idx.append(src_i)
|
||||
kept.append((x0, y0, x1, y1, pelv))
|
||||
if len(keep_idx) >= self.num_persons:
|
||||
break
|
||||
humans = [humans[i] for i in keep_idx]
|
||||
n_keep = len(humans)
|
||||
|
||||
bboxes = []
|
||||
for hh in humans:
|
||||
v = hh["v3d"].detach().cpu().numpy()
|
||||
xmin, ymin = float(v[:, 0].min()), float(v[:, 1].min())
|
||||
xmax, ymax = float(v[:, 0].max()), float(v[:, 1].max())
|
||||
bboxes.append([PoseKp(x=xmin, y=ymin, c=1.0),
|
||||
PoseKp(x=xmax, y=ymax, c=1.0)])
|
||||
ids = self._tracker.update(bboxes)
|
||||
|
||||
persons: list[SMPLXPerson] = []
|
||||
for i, hh in enumerate(humans[:n_keep]):
|
||||
pid = ids[i] if i < len(ids) else i
|
||||
if pid < 0:
|
||||
continue
|
||||
v3d = hh["v3d"].detach().cpu().numpy()
|
||||
transl_np = hh["transl_pelvis"].detach().cpu().numpy().flatten()
|
||||
shape_raw = hh["shape"].detach().cpu().numpy().flatten()
|
||||
expr_raw = hh["expression"].detach().cpu().numpy().flatten()
|
||||
|
||||
pid_c = pid % self.num_persons
|
||||
shape_n = min(10, len(shape_raw))
|
||||
expr_n = min(10, len(expr_raw))
|
||||
shape_smooth = np.zeros(10, dtype=np.float32)
|
||||
expr_smooth = np.zeros(10, dtype=np.float32)
|
||||
for k in range(shape_n):
|
||||
shape_smooth[k] = self._smooth_shape[pid_c][k](
|
||||
float(shape_raw[k]), t_now)
|
||||
for k in range(expr_n):
|
||||
expr_smooth[k] = self._smooth_expr[pid_c][k](
|
||||
float(expr_raw[k]), t_now)
|
||||
|
||||
persons.append(SMPLXPerson(
|
||||
pid=int(pid),
|
||||
vertices_3d=np.ascontiguousarray(v3d, dtype=np.float32),
|
||||
translation=np.ascontiguousarray(
|
||||
transl_np[:3], dtype=np.float32),
|
||||
confidence=float(hh["scores"].item()),
|
||||
betas=np.ascontiguousarray(shape_smooth, dtype=np.float32),
|
||||
expression=np.ascontiguousarray(expr_smooth, dtype=np.float32),
|
||||
))
|
||||
|
||||
with self.state.lock():
|
||||
self.state.persons_smplx = persons
|
||||
self.state.smplx_last_t = t_now
|
||||
|
||||
t_end = time.monotonic()
|
||||
dt_total = (t_end - t_cap_start) * 1e3
|
||||
if LOG.isEnabledFor(logging.DEBUG) or dt_total > 100.0:
|
||||
LOG.log(
|
||||
logging.DEBUG if dt_total <= 100.0 else logging.WARNING,
|
||||
"frame[coreml]: cap=%.1f pre=%.1f inf=%.1f "
|
||||
"post=%.1fms total=%.1fms",
|
||||
(t_pre_start - t_cap_start) * 1e3,
|
||||
(t_inf_start - t_pre_start) * 1e3,
|
||||
(t_post_start - t_inf_start) * 1e3,
|
||||
(t_end - t_post_start) * 1e3,
|
||||
dt_total,
|
||||
)
|
||||
|
||||
dt = time.monotonic() - t_cap_start
|
||||
if dt < self.period:
|
||||
time.sleep(self.period - dt)
|
||||
|
||||
cap.stop()
|
||||
LOG.info("multi_hmr coreml worker stopped")
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
"""Multi-HMR CoreML backend (ANE/GPU/CPU via Apple's CoreML framework).
|
||||
|
||||
Python 3.14 cannot use `coremltools.MLModel` because `libcoremlpython`
|
||||
and `libmilstoragepython` native extensions are not distributed for
|
||||
3.14. We load CoreML.framework directly via `objc.loadBundle()` —
|
||||
same pattern as `coreml_pose.py`.
|
||||
|
||||
Unlike `coreml_pose.py`, this backend does NOT use Vision: Vision is
|
||||
limited to image inputs and cannot feed a second MLMultiArray (cam_K).
|
||||
We invoke `MLModel.predictionFromFeatures:error:` directly with a
|
||||
`MLDictionaryFeatureProvider` wrapping two `MLMultiArray`s.
|
||||
|
||||
Public API:
|
||||
backend = MultiHMRCoreMLBackend(mlpackage_path)
|
||||
humans = backend.infer(image_chw_f32, K_33_f32, det_thresh=0.3)
|
||||
# humans is a list[dict] with the same keys as the PyTorch model
|
||||
# output. Values are CoreMLArray instances that quack like torch
|
||||
# tensors (.detach().cpu().numpy() / .item()).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import objc
|
||||
from Foundation import NSURL
|
||||
|
||||
LOG = logging.getLogger("multihmr_coreml")
|
||||
|
||||
DEFAULT_MLPACKAGE = (
|
||||
Path.home() / ".cache" / "av-live-multihmr"
|
||||
/ "multihmr_full_672_s.mlpackage"
|
||||
)
|
||||
|
||||
# Multi-HMR exported with apply_topk(K=4): outputs are fixed shape.
|
||||
N_PERSONS_FIXED = 4
|
||||
N_VERTS = 10475
|
||||
|
||||
# CoreML output names from the exported .mlpackage.
|
||||
OUT_V3D = "var_2541" # (4, 10475, 3) f16
|
||||
OUT_TRANSL = "var_2544" # (4, 1, 3) f16
|
||||
OUT_SCORES = "var_2557" # (4,) f16
|
||||
OUT_BETAS = "var_2560" # (4, 10) f16
|
||||
OUT_EXPR = "var_2563" # (4, 10) f16
|
||||
|
||||
# MLMultiArrayDataType raw values (from CoreML headers).
|
||||
ML_DTYPE_FLOAT32 = 65568
|
||||
ML_DTYPE_FLOAT16 = 65552
|
||||
ML_DTYPE_DOUBLE = 65600
|
||||
ML_DTYPE_INT32 = 131104
|
||||
|
||||
|
||||
_NS: dict[str, Any] = {}
|
||||
_FRAMEWORKS_LOADED = False
|
||||
|
||||
|
||||
def _load_frameworks() -> dict[str, Any]:
|
||||
global _FRAMEWORKS_LOADED
|
||||
if _FRAMEWORKS_LOADED:
|
||||
return _NS
|
||||
objc.loadBundle("CoreML", _NS,
|
||||
"/System/Library/Frameworks/CoreML.framework")
|
||||
_FRAMEWORKS_LOADED = True
|
||||
return _NS
|
||||
|
||||
|
||||
class CoreMLArray:
|
||||
"""Tiny tensor-like adapter so the existing worker hot path can
|
||||
treat CoreML outputs the same way it treats torch tensors.
|
||||
|
||||
Supports `.detach().cpu().numpy()` and `.item()`. The wrapper is
|
||||
a no-op around a numpy array; we keep the chain so callers don't
|
||||
need any conditional branch."""
|
||||
|
||||
__slots__ = ("_arr",)
|
||||
|
||||
def __init__(self, arr: np.ndarray) -> None:
|
||||
self._arr = arr
|
||||
|
||||
def detach(self) -> "CoreMLArray":
|
||||
return self
|
||||
|
||||
def cpu(self) -> "CoreMLArray":
|
||||
return self
|
||||
|
||||
def numpy(self) -> np.ndarray:
|
||||
return self._arr
|
||||
|
||||
def item(self) -> float:
|
||||
return float(self._arr.reshape(-1)[0])
|
||||
|
||||
@property
|
||||
def shape(self) -> tuple[int, ...]:
|
||||
return tuple(self._arr.shape)
|
||||
|
||||
|
||||
def _np_to_mlarray(arr: np.ndarray):
|
||||
"""Create a contiguous float32 MLMultiArray from a numpy array.
|
||||
|
||||
We always feed FLOAT32 — even though outputs are FLOAT16, CoreML
|
||||
will auto-cast on the input side."""
|
||||
ns = _load_frameworks()
|
||||
MLMultiArray = 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")
|
||||
# Copy bytes through dataPointer (raw void*). pyobjc exposes it as
|
||||
# a memoryview-like opaque; we use ctypes to memcpy.
|
||||
import ctypes
|
||||
ptr = ml.dataPointer()
|
||||
n_bytes = arr.nbytes
|
||||
# pyobjc returns either an objc.varlist or a Python int pointer.
|
||||
addr = int(ptr) if isinstance(ptr, int) else ctypes.cast(
|
||||
ptr, ctypes.c_void_p).value
|
||||
if addr is None:
|
||||
raise RuntimeError("MLMultiArray dataPointer null")
|
||||
ctypes.memmove(addr, arr.ctypes.data, n_bytes)
|
||||
return ml
|
||||
|
||||
|
||||
def _mlarray_to_np(ml) -> np.ndarray:
|
||||
"""Copy an MLMultiArray (FLOAT16 or FLOAT32) into a numpy float32."""
|
||||
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("MLMultiArray 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 MLMultiArray dtype {dtype_id}")
|
||||
return arr.reshape(shape)
|
||||
|
||||
|
||||
class MultiHMRCoreMLBackend:
|
||||
"""CoreML inference wrapper for Multi-HMR (full_672_s)."""
|
||||
|
||||
def __init__(self, mlpackage_path: Path | 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}")
|
||||
ns = _load_frameworks()
|
||||
MLModel = ns["MLModel"]
|
||||
MLModelConfiguration = ns["MLModelConfiguration"]
|
||||
cfg = MLModelConfiguration.alloc().init()
|
||||
try:
|
||||
# MLComputeUnits: 0=CPUOnly, 1=CPUAndGPU, 2=All (ANE+GPU+CPU),
|
||||
# 3=CPUAndNeuralEngine. Multi-HMR's ANEF compile fails
|
||||
# (validated 2026-05-13 on M5), and 'All' falls back to a
|
||||
# slow path (~146ms). CPU+GPU = 28ms = ~35fps on M5.
|
||||
cfg.setComputeUnits_(1)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
url = NSURL.fileURLWithPath_(str(self.path))
|
||||
# .mlpackage must be compiled to .mlmodelc before MLModel can
|
||||
# load it. compileModelAtURL_error_ returns an NSURL to a
|
||||
# temp .mlmodelc bundle.
|
||||
compiled_url = MLModel.compileModelAtURL_error_(url, None)
|
||||
if compiled_url is None:
|
||||
raise RuntimeError(f"compileModelAtURL failed for {self.path}")
|
||||
model = MLModel.modelWithContentsOfURL_configuration_error_(
|
||||
compiled_url, cfg, None)
|
||||
if model is None:
|
||||
raise RuntimeError(f"MLModel load failed for {compiled_url}")
|
||||
self._model = model
|
||||
self._ns = ns
|
||||
LOG.info("Multi-HMR CoreML model loaded (%s, computeUnits=CPU+GPU)",
|
||||
self.path.name)
|
||||
|
||||
@staticmethod
|
||||
def is_available(mlpackage_path: Path | None = None) -> bool:
|
||||
p = Path(mlpackage_path) if mlpackage_path else DEFAULT_MLPACKAGE
|
||||
if not p.exists():
|
||||
return False
|
||||
try:
|
||||
_load_frameworks()
|
||||
return True
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
|
||||
def _predict(self, image_4d: np.ndarray, K_33: np.ndarray) -> dict:
|
||||
ns = self._ns
|
||||
MLDictionaryFeatureProvider = ns["MLDictionaryFeatureProvider"]
|
||||
MLFeatureValue = ns["MLFeatureValue"]
|
||||
img_ml = _np_to_mlarray(image_4d)
|
||||
k_ml = _np_to_mlarray(K_33)
|
||||
feats = {
|
||||
"image": MLFeatureValue.featureValueWithMultiArray_(img_ml),
|
||||
"cam_K": MLFeatureValue.featureValueWithMultiArray_(k_ml),
|
||||
}
|
||||
provider = MLDictionaryFeatureProvider.alloc(
|
||||
).initWithDictionary_error_(feats, None)
|
||||
if provider is None:
|
||||
raise RuntimeError("MLDictionaryFeatureProvider alloc failed")
|
||||
out = self._model.predictionFromFeatures_error_(provider, None)
|
||||
if out is None:
|
||||
raise RuntimeError("MLModel predict failed")
|
||||
names = [str(n) for n in out.featureNames()]
|
||||
result = {}
|
||||
for n in names:
|
||||
fv = out.featureValueForName_(n)
|
||||
ml = fv.multiArrayValue()
|
||||
if ml is None:
|
||||
continue
|
||||
result[n] = _mlarray_to_np(ml)
|
||||
return result
|
||||
|
||||
def infer(
|
||||
self,
|
||||
image_chw_float32: np.ndarray,
|
||||
K_33: np.ndarray,
|
||||
det_thresh: float = 0.3,
|
||||
) -> list[dict]:
|
||||
"""Run a forward pass and return list of humans dicts.
|
||||
|
||||
Args:
|
||||
image_chw_float32: (3, 672, 672) or (1, 3, 672, 672) in [0,1].
|
||||
K_33: (3, 3) or (1, 3, 3) camera intrinsics.
|
||||
det_thresh: scores threshold; CoreML forwards K=4 always.
|
||||
|
||||
Returns:
|
||||
list[dict] with keys v3d, transl_pelvis, scores, shape,
|
||||
expression. Values are CoreMLArray wrappers.
|
||||
"""
|
||||
img = np.asarray(image_chw_float32, dtype=np.float32)
|
||||
if img.ndim == 3:
|
||||
img = img[np.newaxis, ...]
|
||||
if img.shape != (1, 3, 672, 672):
|
||||
raise ValueError(f"image shape {img.shape}, expected (1,3,672,672)")
|
||||
K = np.asarray(K_33, dtype=np.float32)
|
||||
if K.ndim == 2:
|
||||
K = K[np.newaxis, ...]
|
||||
if K.shape != (1, 3, 3):
|
||||
raise ValueError(f"K shape {K.shape}, expected (1,3,3)")
|
||||
|
||||
raw = self._predict(img, K)
|
||||
v3d = raw.get(OUT_V3D)
|
||||
transl = raw.get(OUT_TRANSL)
|
||||
scores = raw.get(OUT_SCORES)
|
||||
betas = raw.get(OUT_BETAS)
|
||||
expr = raw.get(OUT_EXPR)
|
||||
if any(x is None for x in (v3d, transl, scores, betas, expr)):
|
||||
raise RuntimeError(
|
||||
"missing outputs; got keys=" + ",".join(raw.keys()))
|
||||
|
||||
humans: list[dict] = []
|
||||
for k in range(N_PERSONS_FIXED):
|
||||
sc = float(scores[k])
|
||||
if sc < det_thresh:
|
||||
continue
|
||||
humans.append({
|
||||
"v3d": CoreMLArray(v3d[k]), # (10475, 3)
|
||||
"transl_pelvis": CoreMLArray(transl[k]), # (1, 3)
|
||||
"scores": CoreMLArray(np.array([sc], dtype=np.float32)),
|
||||
"shape": CoreMLArray(betas[k]), # (10,)
|
||||
"expression": CoreMLArray(expr[k]), # (10,)
|
||||
})
|
||||
return humans
|
||||
@@ -170,11 +170,28 @@ _K_INV_PRE = torch.tensor([
|
||||
])
|
||||
|
||||
def inverse_perspective_projection_fixed(points, K, distance):
|
||||
"""Bypass torch.inverse : utilise K_inv pre-calcule en closed-form
|
||||
(notre K est connu et fixe). Le K argument est ignore."""
|
||||
K_inv = _K_INV_PRE.to(points.device).to(points.dtype)
|
||||
points = torch.cat([points, torch.ones_like(points[..., :1])], -1)
|
||||
points = torch.einsum('bij,bkj->bki', K_inv, points)
|
||||
"""Bypass torch.inverse + einsum + matmul pour eviter le bug
|
||||
coremltools de broadcast batch 1->K sur ces ops. K_inv etant
|
||||
fixe et structure (diag + translate), on ecrit les composantes
|
||||
explicitement en ops elementaires.
|
||||
|
||||
K_inv = [[1/f, 0, -cx/f], [0, 1/f, -cy/f], [0, 0, 1]]
|
||||
Pour points (b, N, 3) : out = points @ K_inv.T donne :
|
||||
out[..., 0] = points[..., 0]/f - (cx/f) * points[..., 2]
|
||||
out[..., 1] = points[..., 1]/f - (cy/f) * points[..., 2]
|
||||
out[..., 2] = points[..., 2]
|
||||
"""
|
||||
points_hom = torch.cat([points, torch.ones_like(points[..., :1])], -1)
|
||||
inv_f = 1.0 / focal_val
|
||||
cx_over_f = cx / focal_val
|
||||
cy_over_f = cy / focal_val
|
||||
x = points_hom[..., 0:1]
|
||||
y = points_hom[..., 1:2]
|
||||
z = points_hom[..., 2:3]
|
||||
out0 = x * inv_f - z * cx_over_f
|
||||
out1 = y * inv_f - z * cy_over_f
|
||||
out2 = z
|
||||
points = torch.cat([out0, out1, out2], dim=-1)
|
||||
if distance is None:
|
||||
return points
|
||||
points = points * distance
|
||||
@@ -190,6 +207,26 @@ model_mod.inverse_perspective_projection = inverse_perspective_projection_fixed
|
||||
import blocks.smpl_layer as _smpl_layer
|
||||
_smpl_layer.inverse_perspective_projection = inverse_perspective_projection_fixed
|
||||
|
||||
# Aussi perspective_projection (utilise dans smpl_layer.py:143-144 pour
|
||||
# j2d et v2d) -> rewrite einsum en matmul pour le meme broadcast bug.
|
||||
def perspective_projection_fixed(x, K):
|
||||
"""Element-wise rewrite de la projection perspective avec K fixe
|
||||
(focal=IMG_SIZE, cx=cy=IMG_SIZE/2). Bypass matmul/einsum pour eviter
|
||||
les bugs broadcast coremltools.
|
||||
K = [[f, 0, cx], [0, f, cy], [0, 0, 1]]
|
||||
out[..., 0] = f * x_norm + cx * z_norm (mais on veut [..., :2])
|
||||
= f * (x/z) + cx
|
||||
out[..., 1] = f * (y/z) + cy
|
||||
"""
|
||||
z = x[..., 2:3]
|
||||
px = x[..., 0:1] / z * focal_val + cx
|
||||
py = x[..., 1:2] / z * focal_val + cy
|
||||
return torch.cat([px, py], dim=-1)
|
||||
|
||||
_camera.perspective_projection = perspective_projection_fixed
|
||||
_utils_pkg.perspective_projection = perspective_projection_fixed
|
||||
_smpl_layer.perspective_projection = perspective_projection_fixed
|
||||
|
||||
|
||||
# === Wrapper qui produit tuple fixe ===
|
||||
class TracedMHMR(nn.Module):
|
||||
@@ -422,6 +459,28 @@ def _diagonal_general(context, node):
|
||||
|
||||
_TORCH_OPS_REGISTRY.name_to_func_mapping["diagonal"] = _diagonal_general
|
||||
|
||||
|
||||
# Instrument reshape pour logger node source au moment de l'erreur.
|
||||
from coremltools.converters.mil.mil.ops.defs.iOS15 import tensor_transformation as _tt
|
||||
_orig_reshape_ti = _tt.reshape.type_inference
|
||||
|
||||
|
||||
def _reshape_ti_logged(self):
|
||||
try:
|
||||
return _orig_reshape_ti(self)
|
||||
except ValueError as e:
|
||||
if "Invalid target shape" in str(e):
|
||||
try:
|
||||
from_shape = list(self.x.shape)
|
||||
target = list(self.shape.val) if hasattr(self.shape, "val") else "?"
|
||||
print(f" >>> RESHAPE FAIL : name={self.name} from={from_shape} target={target}")
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
_tt.reshape.type_inference = _reshape_ti_logged
|
||||
|
||||
try:
|
||||
mlmodel = ct.convert(
|
||||
traced,
|
||||
|
||||
@@ -49,4 +49,71 @@ if [ ! -e "$CACHE/multi-hmr/models" ]; then
|
||||
ln -sfn ../models "$CACHE/multi-hmr/models"
|
||||
fi
|
||||
|
||||
# CoreML conversion patches : remplace les torch.einsum dans utils/camera.py
|
||||
# par des ops element-wise (broadcast-friendly). Sans ca, ct.convert echoue
|
||||
# avec "Invalid target shape in reshape op ([1, N, 3] to [K*N, 3, 1])"
|
||||
# quand batch K detections != 1. Idempotent.
|
||||
CAM="$CACHE/multi-hmr/utils/camera.py"
|
||||
if [ -f "$CAM" ] && ! grep -q "_apply_intrinsics_componentwise" "$CAM"; then
|
||||
echo "==> Patch utils/camera.py (einsum -> componentwise)"
|
||||
python3 - "$CAM" <<'PYEOF'
|
||||
import sys, pathlib
|
||||
p = pathlib.Path(sys.argv[1])
|
||||
src = p.read_text()
|
||||
helper = '''
|
||||
def _apply_intrinsics_componentwise(K, y):
|
||||
"""CoreML-friendly: out[b,k,i] = sum_j K[b,i,j] * y[b,k,j]
|
||||
Replaces torch.einsum('bij,bkj->bki', K, y) with pure broadcast ops.
|
||||
"""
|
||||
K00 = K[:, 0:1, 0:1]; K01 = K[:, 0:1, 1:2]; K02 = K[:, 0:1, 2:3]
|
||||
K10 = K[:, 1:2, 0:1]; K11 = K[:, 1:2, 1:2]; K12 = K[:, 1:2, 2:3]
|
||||
K20 = K[:, 2:3, 0:1]; K21 = K[:, 2:3, 1:2]; K22 = K[:, 2:3, 2:3]
|
||||
y0 = y[:, :, 0:1]; y1 = y[:, :, 1:2]; y2 = y[:, :, 2:3]
|
||||
out0 = K00 * y0 + K01 * y1 + K02 * y2
|
||||
out1 = K10 * y0 + K11 * y1 + K12 * y2
|
||||
out2 = K20 * y0 + K21 * y1 + K22 * y2
|
||||
return torch.cat([out0, out1, out2], dim=-1)
|
||||
|
||||
|
||||
'''
|
||||
src = src.replace(
|
||||
"def perspective_projection(x, K):",
|
||||
helper + "def perspective_projection(x, K):",
|
||||
)
|
||||
src = src.replace(
|
||||
"y = torch.einsum('bij,bkj->bki', K, y) # (bs, N, 3)",
|
||||
"y = _apply_intrinsics_componentwise(K, y)",
|
||||
)
|
||||
src = src.replace(
|
||||
"points = torch.einsum('bij,bkj->bki', torch.inverse(K), points)",
|
||||
"points = _apply_intrinsics_componentwise(torch.inverse(K), points)",
|
||||
)
|
||||
p.write_text(src)
|
||||
print(" camera.py patched")
|
||||
PYEOF
|
||||
fi
|
||||
|
||||
# CoreML conversion patch : smplx/lbs.py landmarks einsum (mеme bug broadcast)
|
||||
# Patch best-effort sur tous les venvs presents (data_only_viz + /tmp/coreml312).
|
||||
for VENV in \
|
||||
"$(dirname "$(dirname "$(readlink -f "$0")")")/.venv" \
|
||||
"/tmp/coreml312"; do
|
||||
LBS="$VENV/lib/python3.14/site-packages/smplx/lbs.py"
|
||||
[ -f "$LBS" ] || LBS="$VENV/lib/python3.12/site-packages/smplx/lbs.py"
|
||||
if [ -f "$LBS" ] && grep -q "torch.einsum('blfi,blf->bli'" "$LBS"; then
|
||||
echo "==> Patch $LBS (landmarks einsum)"
|
||||
python3 - "$LBS" <<'PYEOF'
|
||||
import sys, pathlib
|
||||
p = pathlib.Path(sys.argv[1])
|
||||
s = p.read_text()
|
||||
s = s.replace(
|
||||
"landmarks = torch.einsum('blfi,blf->bli', [lmk_vertices, lmk_bary_coords])\n return landmarks",
|
||||
"# CoreML-friendly: replace einsum('blfi,blf->bli', ...) with broadcast+sum\n landmarks = (lmk_vertices * lmk_bary_coords.unsqueeze(-1)).sum(dim=2)\n return landmarks",
|
||||
)
|
||||
p.write_text(s)
|
||||
print(" smplx/lbs.py patched")
|
||||
PYEOF
|
||||
fi
|
||||
done
|
||||
|
||||
echo "Setup OK. Cache : $CACHE"
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Tests for the Multi-HMR CoreML backend.
|
||||
|
||||
Skipped unless the .mlpackage exists at the standard cache path.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
MLPACKAGE = (
|
||||
Path.home() / ".cache" / "av-live-multihmr"
|
||||
/ "multihmr_full_672_s.mlpackage"
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not MLPACKAGE.exists(),
|
||||
reason=f"mlpackage missing at {MLPACKAGE}",
|
||||
)
|
||||
|
||||
|
||||
def _make_K() -> np.ndarray:
|
||||
f = 672.0
|
||||
return np.array([[f, 0.0, 336.0],
|
||||
[0.0, f, 336.0],
|
||||
[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
|
||||
def test_is_available_true():
|
||||
from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend
|
||||
assert MultiHMRCoreMLBackend.is_available(MLPACKAGE) is True
|
||||
|
||||
|
||||
def test_load_model():
|
||||
from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend
|
||||
backend = MultiHMRCoreMLBackend(MLPACKAGE)
|
||||
assert backend._model is not None
|
||||
|
||||
|
||||
def test_infer_random_image_shapes():
|
||||
from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend
|
||||
backend = MultiHMRCoreMLBackend(MLPACKAGE)
|
||||
rng = np.random.default_rng(0)
|
||||
img = rng.random((3, 672, 672), dtype=np.float32)
|
||||
K = _make_K()
|
||||
# threshold = -inf so we get all K=4 humans back
|
||||
humans = backend.infer(img, K, det_thresh=-1.0)
|
||||
assert len(humans) == 4
|
||||
for h in humans:
|
||||
v = h["v3d"].detach().cpu().numpy()
|
||||
assert v.shape == (10475, 3)
|
||||
assert v.dtype == np.float32
|
||||
t = h["transl_pelvis"].detach().cpu().numpy()
|
||||
assert t.shape == (1, 3)
|
||||
s = float(h["scores"].item())
|
||||
assert isinstance(s, float)
|
||||
beta = h["shape"].detach().cpu().numpy()
|
||||
assert beta.shape == (10,)
|
||||
expr = h["expression"].detach().cpu().numpy()
|
||||
assert expr.shape == (10,)
|
||||
|
||||
|
||||
def test_infer_latency_under_target():
|
||||
from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend
|
||||
backend = MultiHMRCoreMLBackend(MLPACKAGE)
|
||||
K = _make_K()
|
||||
rng = np.random.default_rng(42)
|
||||
img = rng.random((3, 672, 672), dtype=np.float32)
|
||||
# warmup
|
||||
backend.infer(img, K, det_thresh=-1.0)
|
||||
# measure
|
||||
n = 5
|
||||
times = []
|
||||
for _ in range(n):
|
||||
t0 = time.monotonic()
|
||||
backend.infer(img, K, det_thresh=-1.0)
|
||||
times.append((time.monotonic() - t0) * 1e3)
|
||||
times.sort()
|
||||
median_ms = times[n // 2]
|
||||
print(f"median latency: {median_ms:.1f} ms (n={n})")
|
||||
# Target 50ms = 20fps. M5 bench shows ~29ms. Generous margin.
|
||||
assert median_ms < 80.0, f"median {median_ms:.1f}ms > 80ms target"
|
||||
|
||||
|
||||
def test_filter_threshold():
|
||||
from data_only_viz.multihmr_coreml import MultiHMRCoreMLBackend
|
||||
backend = MultiHMRCoreMLBackend(MLPACKAGE)
|
||||
rng = np.random.default_rng(0)
|
||||
img = rng.random((3, 672, 672), dtype=np.float32)
|
||||
K = _make_K()
|
||||
high = backend.infer(img, K, det_thresh=999.0)
|
||||
assert high == [] # nothing passes
|
||||
low = backend.infer(img, K, det_thresh=-1.0)
|
||||
assert len(low) == 4
|
||||
Reference in New Issue
Block a user