diff --git a/data_only_viz/scripts/coreml_full_probe.py b/data_only_viz/scripts/coreml_full_probe.py new file mode 100644 index 0000000..e55ca4f --- /dev/null +++ b/data_only_viz/scripts/coreml_full_probe.py @@ -0,0 +1,251 @@ +"""Task 3 — Convert FULL Multi-HMR (backbone + head) to CoreML +avec apply_topk(K=4) + fixed-shape tuple output. +""" +from __future__ import annotations + +import os +import sys +import time +import types +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn + +CACHE = Path.home() / ".cache" / "av-live-multihmr" +CKPT = CACHE / "checkpoints" / "multiHMR_672_S.pt" +MULTIHMR_REPO = CACHE / "multi-hmr" + +sys.path.insert(0, str(MULTIHMR_REPO)) +for mod in ("pyrender", "pyvista", "anny"): + sys.modules.setdefault(mod, types.ModuleType(mod)) + +DEVICE = "cpu" # trace on CPU (CoreML mlprogram doesn't care) +IMG_SIZE = 672 +K_PERSONS = 4 + + +# === apply_topk replacement (validated equivalent in Task 2) === +def apply_topk(K, _scores): + if isinstance(K, list): + K = K[0] + B, H, W, C = _scores.shape + flat = _scores.reshape(B, -1) + _, idx_flat = torch.topk(flat, k=K, dim=1) + wc = W * C + idx_b = (torch.arange(B, device=_scores.device) + .unsqueeze(1).expand(-1, K).reshape(-1).long()) + idx_flat_flat = idx_flat.reshape(-1) + idx_h = (idx_flat_flat // wc).long() + idx_w = ((idx_flat_flat // C) % W).long() + idx_c = (idx_flat_flat % C).long() + return (idx_b, idx_h, idx_w, idx_c) + + +# === Patch coremltools _cast (validated probe v4) === +def _patched_cast(context, node, dtype, dtype_str): + from coremltools.converters.mil import Builder as mb + from coremltools.converters.mil.frontend.torch import ops as _ops + inputs = _ops._get_inputs(context, node, expected=1) + x = inputs[0] + if x.val is not None: + try: + const_val = dtype(x.val) + except TypeError: + arr = np.asarray(x.val) + if arr.size == 1: + const_val = dtype(arr.item()) + else: + res = mb.cast(x=x, dtype=dtype_str, name=node.name) + context.add(res) + return + res = mb.const(val=const_val, name=node.name) + else: + res = mb.cast(x=x, dtype=dtype_str, name=node.name) + context.add(res) + + +prev = os.getcwd() +try: + os.chdir(MULTIHMR_REPO) + from model import Model + import model as model_mod + + # Inject topk replacement + print("==> Patching apply_threshold -> apply_topk(K=4)") + model_mod.apply_threshold = lambda thr, scores: apply_topk(K_PERSONS, scores) + + torch_dev = torch.device(DEVICE) + ckpt = torch.load(str(CKPT), map_location=torch_dev, weights_only=False) + kw = {k: v for k, v in vars(ckpt["args"]).items()} + kw["type"] = ckpt["args"].train_return_type + kw["img_size"] = ckpt["args"].img_size[0] + print(f"==> Loading Multi-HMR ViT-S 672 (params count tbd)") + model = Model(**kw).to(torch_dev) + model.load_state_dict(ckpt["model_state_dict"], strict=False) + model.eval() +finally: + os.chdir(prev) + + +# === Pre-compute interpolate_pos_encoding (probe v4 fix) === +# Multi-HMR's backbone is DINOv2 ViT-S/14 — same dynamic interpolation +# problem that planted la conversion sur backbone seul. +if hasattr(model.backbone, "encoder") and hasattr(model.backbone.encoder, + "interpolate_pos_encoding"): + print("==> Patching backbone.encoder.interpolate_pos_encoding (pre-compute)") + bk = model.backbone.encoder + with torch.no_grad(): + dummy_x = torch.rand(1, 3, IMG_SIZE, IMG_SIZE) + dummy_p = bk.patch_embed(dummy_x) + cls = bk.cls_token.expand(dummy_p.shape[0], -1, -1) + x_full = torch.cat((cls, dummy_p), dim=1) + cached_pe = bk.interpolate_pos_encoding( + x_full, IMG_SIZE, IMG_SIZE).detach() + bk.register_buffer("_cached_pos_embed", cached_pe) + + def fixed_pe(self, x, w, h): + return self._cached_pos_embed.to(x.dtype) + bk.interpolate_pos_encoding = types.MethodType(fixed_pe, bk) + print(f" cached shape {tuple(cached_pe.shape)}") + +# === Patch utils.camera.inverse_perspective_projection === +# torch.inverse(K) plante coremltools (op non implementee). Comme K est +# fixe (camera intrinsics avec focal=IMG_SIZE), on pre-calcule K_inv +# en closed-form et on l'utilise comme buffer module-level. +print("==> Patching utils.camera.inverse_perspective_projection") +import utils.camera as _camera + +# Pre-compute K_inv closed-form pour notre K standard +focal_val = float(IMG_SIZE) +cx = cy = IMG_SIZE / 2.0 +_K_INV_PRE = torch.tensor([ + [[1.0 / focal_val, 0.0, -cx / focal_val], + [0.0, 1.0 / focal_val, -cy / focal_val], + [0.0, 0.0, 1.0]] +]) + +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) + if distance is None: + return points + points = points * distance + return points + +_camera.inverse_perspective_projection = inverse_perspective_projection_fixed +# Aussi patcher le re-export dans utils/__init__.py et model.py +import utils as _utils_pkg +_utils_pkg.inverse_perspective_projection = inverse_perspective_projection_fixed +# model.py importe directement : monkey-patch sur le module +model_mod.inverse_perspective_projection = inverse_perspective_projection_fixed +# Idem smpl_layer +import blocks.smpl_layer as _smpl_layer +_smpl_layer.inverse_perspective_projection = inverse_perspective_projection_fixed + + +# === Wrapper qui produit tuple fixe === +class TracedMHMR(nn.Module): + """Wrap Multi-HMR pour trace : output tuple de tensors fixes, + pas de list-of-dicts.""" + + def __init__(self, m: Model): + super().__init__() + self.m = m + + def forward(self, x: torch.Tensor, cam_K: torch.Tensor): + # Call original forward with is_training=False ; apply_topk + # garantit toujours K=4 detections donc le loop dans le forward + # est unroll-friendly. + humans = self.m(x, is_training=False, nms_kernel_size=5, + det_thresh=0.0, K=cam_K) + # humans est une list[dict] de longueur 4. Stack en tensors. + if len(humans) == 0: + # Should not happen with apply_topk, but defensive + zeros = torch.zeros(K_PERSONS, 10475, 3) + zeros_p = torch.zeros(K_PERSONS, 3) + zeros_s = torch.zeros(K_PERSONS) + zeros_b = torch.zeros(K_PERSONS, 10) + return zeros, zeros_p, zeros_s, zeros_b, zeros_b + v3d = torch.stack([h["v3d"] for h in humans]) + transl = torch.stack([h["transl_pelvis"] for h in humans]) + scores = torch.stack([ + h["scores"] if h["scores"].dim() > 0 else h["scores"].unsqueeze(0) + for h in humans + ]).squeeze(-1) + shape = torch.stack([h["shape"] for h in humans]) + expr = torch.stack([h["expression"] for h in humans]) + return v3d, transl, scores, shape, expr + + +wrapper = TracedMHMR(model).eval() + +# Sanity forward +focal = float(IMG_SIZE) +example_K = torch.tensor( + [[[focal, 0.0, IMG_SIZE / 2.0], + [0.0, focal, IMG_SIZE / 2.0], + [0.0, 0.0, 1.0]]], dtype=torch.float32) +example_x = torch.rand(1, 3, IMG_SIZE, IMG_SIZE) + +print("==> Sanity forward") +with torch.no_grad(): + v3d, transl, scores, shape, expr = wrapper(example_x, example_K) +print(f" v3d: {tuple(v3d.shape)}, transl: {tuple(transl.shape)},") +print(f" scores: {tuple(scores.shape)}, shape: {tuple(shape.shape)},") +print(f" expr: {tuple(expr.shape)}") + +print("==> torch.jit.trace") +try: + traced = torch.jit.trace(wrapper, (example_x, example_K), strict=False) + print(" trace OK") +except Exception as e: + print(f" trace FAILED: {type(e).__name__}: {e}") + raise + +# === CoreML convert === +print("==> coremltools.convert") +import coremltools as ct +from coremltools.converters.mil.frontend.torch import ops as _ops +_ops._cast = _patched_cast + +try: + mlmodel = ct.convert( + traced, + inputs=[ + ct.TensorType(shape=(1, 3, IMG_SIZE, IMG_SIZE), + name="image", dtype=np.float32), + ct.TensorType(shape=(1, 3, 3), name="cam_K", dtype=np.float32), + ], + compute_units=ct.ComputeUnit.CPU_AND_GPU, + minimum_deployment_target=ct.target.macOS15, + convert_to="mlprogram", + ) + out_path = "/tmp/multihmr_full_672_s.mlpackage" + mlmodel.save(out_path) + print(f" CONVERT OK -> {out_path}") +except Exception as e: + print(f" CONVERT FAILED: {type(e).__name__}: {e}") + raise + +# === Bench === +print("==> bench 30 iter") +img = np.random.rand(1, 3, IMG_SIZE, IMG_SIZE).astype(np.float32) +cam = np.array([[[focal, 0, IMG_SIZE/2], + [0, focal, IMG_SIZE/2], + [0, 0, 1]]], dtype=np.float32) +for _ in range(3): + _ = mlmodel.predict({"image": img, "cam_K": cam}) +t = [] +for _ in range(30): + t0 = time.perf_counter() + _ = mlmodel.predict({"image": img, "cam_K": cam}) + t.append((time.perf_counter() - t0) * 1000) +t.sort() +print(f" CoreML full Multi-HMR median={t[15]:.1f} ms " + f"p10={t[3]:.1f} p90={t[27]:.1f} min={t[0]:.1f}") +print(f" Target was <60ms (12-25 fps). Achieved: {1000.0/t[15]:.1f} fps") diff --git a/data_only_viz/scripts/probe_head_topk.py b/data_only_viz/scripts/probe_head_topk.py new file mode 100644 index 0000000..dc103f6 --- /dev/null +++ b/data_only_viz/scripts/probe_head_topk.py @@ -0,0 +1,145 @@ +"""Task 2 — Validate apply_topk(K=4) as drop-in replacement for +apply_threshold in Multi-HMR head. + +Compares v3d output between threshold-based (original) and topk-based +(patched) Multi-HMR on the same input. Pass criterion: for the same +detections (when K >= n_threshold_detected), v3d cosine similarity > 0.99. +""" +from __future__ import annotations + +import os +import sys +import types +from pathlib import Path + +import numpy as np +import torch + +CACHE = Path.home() / ".cache" / "av-live-multihmr" +CKPT = CACHE / "checkpoints" / "multiHMR_672_S.pt" +MULTIHMR_REPO = CACHE / "multi-hmr" + +sys.path.insert(0, str(MULTIHMR_REPO)) +for mod in ("pyrender", "pyvista", "anny"): + sys.modules.setdefault(mod, types.ModuleType(mod)) + +DEVICE = "mps" if torch.backends.mps.is_available() else "cpu" +IMG_SIZE = 672 + + +def apply_topk(K, _scores): + """Drop-in pour apply_threshold. _scores shape (B, H, W, C). + Renvoie 4-tuple LongTensor (batch_idx, h_idx, w_idx, c_idx) chacun + de longueur B*K (au lieu de variable). K candidate top-scoring + tokens par image. + """ + if isinstance(K, list): + K = K[0] + B, H, W, C = _scores.shape + flat = _scores.reshape(B, -1) + _, idx_flat = torch.topk(flat, k=K, dim=1) + wc = W * C + idx_b = (torch.arange(B, device=_scores.device) + .unsqueeze(1).expand(-1, K).reshape(-1)) + idx_flat_flat = idx_flat.reshape(-1) + idx_h = idx_flat_flat // wc + idx_w = (idx_flat_flat // C) % W + idx_c = idx_flat_flat % C + return (idx_b.long(), idx_h.long(), idx_w.long(), idx_c.long()) + + +prev = os.getcwd() +try: + os.chdir(MULTIHMR_REPO) + from model import Model + import model as model_mod + torch_dev = torch.device(DEVICE) + ckpt = torch.load(str(CKPT), map_location=torch_dev, weights_only=False) + kw = {k: v for k, v in vars(ckpt["args"]).items()} + kw["type"] = ckpt["args"].train_return_type + kw["img_size"] = ckpt["args"].img_size[0] + model = Model(**kw).to(torch_dev) + model.load_state_dict(ckpt["model_state_dict"], strict=False) + model.eval() +finally: + os.chdir(prev) + +focal = float(IMG_SIZE) +K_mat = torch.tensor([[[focal, 0.0, IMG_SIZE / 2.0], + [0.0, focal, IMG_SIZE / 2.0], + [0.0, 0.0, 1.0]]], device=DEVICE) + +# Use a real test image (multi-hmr example or webcam capture) +import cv2 +img_path = "/Users/electron/.cache/av-live-multihmr/multi-hmr/example_data/4446582661_b188f82f3c_c.jpg" +img = cv2.imread(img_path) +h, w = img.shape[:2] +side = min(h, w) +y0 = (h - side) // 2; x0 = (w - side) // 2 +img = cv2.resize(img[y0:y0+side, x0:x0+side], (IMG_SIZE, IMG_SIZE)) +img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) +x = (torch.from_numpy(img_rgb).permute(2, 0, 1).float() / 255.0 + ).unsqueeze(0).to(DEVICE) +print(f"loaded image {img_path}") + +# --- Pass 1 : original apply_threshold a tres bas seuil --- +print("==> Pass 1 : apply_threshold(0.15)") +with torch.no_grad(): + humans_orig = model(x, is_training=False, nms_kernel_size=5, + det_thresh=0.15, K=K_mat) +print(f" detected: {len(humans_orig)}") +for i, h in enumerate(humans_orig[:4]): + sc = h.get("scores", 0.0) + if hasattr(sc, "item"): + sc = sc.item() + print(f" [{i}] score={sc:.3f} v3d.shape={tuple(h['v3d'].shape)}") + +# --- Pass 2 : monkey-patch apply_threshold avec apply_topk(K=4) --- +print("\n==> Pass 2 : apply_topk(K=4)") +original_apply_threshold = model_mod.apply_threshold + +def topk_wrapper(det_thresh, _scores): + return apply_topk(4, _scores) + +model_mod.apply_threshold = topk_wrapper + +with torch.no_grad(): + humans_topk = model(x, is_training=False, nms_kernel_size=5, + det_thresh=0.15, K=K_mat) +print(f" detected: {len(humans_topk)}") +for i, h in enumerate(humans_topk[:4]): + sc = h.get("scores", 0.0) + if hasattr(sc, "item"): + sc = sc.item() + print(f" [{i}] score={sc:.3f} v3d.shape={tuple(h['v3d'].shape)}") + +# Restore +model_mod.apply_threshold = original_apply_threshold + +# --- Comparison --- +print("\n==> Comparison") +if len(humans_orig) == 0 or len(humans_topk) == 0: + print(" NO DETECTIONS in one path — adjust threshold lower") + sys.exit(0) + +# Match by score (highest first in both) +o = sorted(humans_orig, key=lambda h: -( + h.get("scores", 0).item() if hasattr(h.get("scores", 0), "item") + else h.get("scores", 0)))[:min(len(humans_orig), 4)] +t = sorted(humans_topk, key=lambda h: -( + h.get("scores", 0).item() if hasattr(h.get("scores", 0), "item") + else h.get("scores", 0)))[:len(o)] + +for i, (ho, ht) in enumerate(zip(o, t)): + vo = ho["v3d"].detach().cpu().numpy().flatten() + vt = ht["v3d"].detach().cpu().numpy().flatten() + dot = float(np.dot(vo, vt)) + nv = float(np.linalg.norm(vo) * np.linalg.norm(vt) + 1e-9) + cos = dot / nv + mae = float(np.mean(np.abs(vo - vt))) + sco = (ho.get("scores", 0).item() + if hasattr(ho.get("scores", 0), "item") else ho.get("scores", 0)) + sct = (ht.get("scores", 0).item() + if hasattr(ht.get("scores", 0), "item") else ht.get("scores", 0)) + print(f" [{i}] cosine={cos:.6f} mae={mae*1000:.3f}mm " + f"score_orig={sco:.4f} score_topk={sct:.4f}") diff --git a/docs/superpowers/plans/2026-05-13-multihmr-coreml-conversion.md b/docs/superpowers/plans/2026-05-13-multihmr-coreml-conversion.md index efda94f..1794a7d 100644 --- a/docs/superpowers/plans/2026-05-13-multihmr-coreml-conversion.md +++ b/docs/superpowers/plans/2026-05-13-multihmr-coreml-conversion.md @@ -67,6 +67,33 @@ these as untraceable. single-day estimate was optimistic — backbone alone needs ~½ day of surgery before head can even be touched. +### Task 2 + 3 attempted (2026-05-13, post-breakthrough) + +**Task 2 — apply_topk validation : PASS** + +`apply_topk(K=4)` validée comme drop-in pour `apply_threshold` : +cosine sim 1.000000 sur 4/4 détections (image example_data), MAE +2-4 mm = bruit float-reorder. Script : `scripts/probe_head_topk.py`. + +**Task 3 — Full Multi-HMR convert : en cours, patches itératifs** + +Tentative conversion full Multi-HMR avec : +- apply_threshold monkey-patched → apply_topk(K=4) +- backbone.encoder.interpolate_pos_encoding → buffer fige +- utils.camera.inverse_perspective_projection → closed-form K_inv + +→ TracedMHMR.forward sanity OK, jit.trace OK, mais conversion +coremltools échoue successivement sur : +1. `upsample_bicubic2d` ✅ résolu par pos_embed fix +2. `aten::inverse` ✅ résolu par closed-form K_inv +3. `Types should have zero-rank ndarray input, got [0.]` ❌ encore + +Pattern : chaque patch débloque la couche suivante. Estimation +restante : 1-2 jours pour couvrir tous les ops résiduels. + +Script Task 3 : `scripts/coreml_full_probe.py` (laissé en place pour +reprise — point d'arrêt précis documenté). + ### Probe v4 (2026-05-13) — BREAKTHROUGH Avec **2 patches au lieu d'1**, la conversion DINOv2 ViT-S 672x672