From 9e7a9f8fd46dc7adbbc46857ff4634b1903a76b2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=27=C3=A9lectron=20rare?= <108685187+electron-rare@users.noreply.github.com> Date: Wed, 13 May 2026 21:42:43 +0200 Subject: [PATCH] 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. --- data_only_viz/multi_hmr_worker.py | 239 ++++++++++++++++- data_only_viz/multihmr_coreml.py | 276 ++++++++++++++++++++ data_only_viz/scripts/coreml_full_probe.py | 69 ++++- data_only_viz/scripts/setup_multihmr.sh | 67 +++++ data_only_viz/tests/test_multihmr_coreml.py | 96 +++++++ 5 files changed, 741 insertions(+), 6 deletions(-) create mode 100644 data_only_viz/multihmr_coreml.py create mode 100644 data_only_viz/tests/test_multihmr_coreml.py diff --git a/data_only_viz/multi_hmr_worker.py b/data_only_viz/multi_hmr_worker.py index 52a64cd..a09092a 100644 --- a/data_only_viz/multi_hmr_worker.py +++ b/data_only_viz/multi_hmr_worker.py @@ -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") diff --git a/data_only_viz/multihmr_coreml.py b/data_only_viz/multihmr_coreml.py new file mode 100644 index 0000000..f46dc4b --- /dev/null +++ b/data_only_viz/multihmr_coreml.py @@ -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 diff --git a/data_only_viz/scripts/coreml_full_probe.py b/data_only_viz/scripts/coreml_full_probe.py index 241ce15..f0e7bfe 100644 --- a/data_only_viz/scripts/coreml_full_probe.py +++ b/data_only_viz/scripts/coreml_full_probe.py @@ -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, diff --git a/data_only_viz/scripts/setup_multihmr.sh b/data_only_viz/scripts/setup_multihmr.sh index b51df54..fcf96cb 100755 --- a/data_only_viz/scripts/setup_multihmr.sh +++ b/data_only_viz/scripts/setup_multihmr.sh @@ -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" diff --git a/data_only_viz/tests/test_multihmr_coreml.py b/data_only_viz/tests/test_multihmr_coreml.py new file mode 100644 index 0000000..2eb7338 --- /dev/null +++ b/data_only_viz/tests/test_multihmr_coreml.py @@ -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