Spaces:
Paused
Paused
| """SAM 3D Body runner (GPU only). Clones facebookresearch/sam-3d-body + detectron2 at first use, | |
| loads the estimator, runs per-frame inference. Heavy imports (torch/cv2/the repo) stay in here so | |
| importing vid2rig never needs them. Designed for a persistent-GPU Space (A10G); needs HF_TOKEN with | |
| gated access to facebook/sam-3d-body-dinov3.""" | |
| import os, sys, subprocess, functools | |
| import numpy as np | |
| SAM3D_DIR = os.environ.get("SAM3D_BODY_ROOT", "/tmp/sam-3d-body") | |
| DETECTRON2 = "git+https://github.com/facebookresearch/detectron2.git@a1ce2f9" | |
| def _pip(*args): | |
| subprocess.run([sys.executable, "-m", "pip", "install", "-q", *args], check=True) | |
| def _ensure_repo(): | |
| if not os.path.isdir(SAM3D_DIR): | |
| subprocess.run(["git", "clone", "--depth", "1", | |
| "https://github.com/facebookresearch/sam-3d-body.git", SAM3D_DIR], check=True) | |
| if SAM3D_DIR not in sys.path: | |
| sys.path.insert(0, SAM3D_DIR) | |
| try: | |
| import detectron2 # noqa: F401 | |
| except Exception: | |
| _pip(DETECTRON2, "--no-build-isolation", "--no-deps") | |
| def _estimator(): | |
| _ensure_repo() | |
| from notebook.utils import setup_sam_3d_body | |
| return setup_sam_3d_body(hf_repo_id="facebook/sam-3d-body-dinov3") | |
| def _log_structure(outs): | |
| print("[sam3d] people detected:", len(outs), flush=True) | |
| if not outs: | |
| return | |
| o = outs[0] | |
| print("[sam3d] output keys:", list(o.keys()), flush=True) | |
| for k, v in o.items(): | |
| shp = getattr(v, "shape", None) | |
| print(f"[sam3d] {k}: shape={tuple(shp) if shp is not None else type(v).__name__}", flush=True) | |
| def _to_list(x): | |
| try: | |
| return x.detach().cpu().numpy().tolist() | |
| except Exception: | |
| try: | |
| return list(x) | |
| except Exception: | |
| return x | |
| def _get_skeleton(est): | |
| """Return the MHR skeleton {names, parents, prerotations, offsets} from the loaded model.""" | |
| out = {"names": None, "parents": None, "prerotations": None, "offsets": None} | |
| try: | |
| char = None | |
| for path in ("model.head_pose.mhr.character_torch", "model.head_pose_hand.mhr.character_torch", | |
| "head_pose.mhr.character_torch"): | |
| cur = est | |
| for p in path.split("."): | |
| cur = getattr(cur, p, None) | |
| if cur is None: | |
| break | |
| if cur is not None: | |
| char = cur; break | |
| if char is None: | |
| print("[sam3d-skel] character not found", flush=True); return out | |
| skel = getattr(char, "skeleton", char) | |
| names = getattr(skel, "joint_names", None) or getattr(char, "joint_names", None) | |
| if names is not None: | |
| out["names"] = [str(n) for n in names] | |
| jp = getattr(skel, "joint_parents", None) | |
| if jp is not None: | |
| out["parents"] = [int(x) for x in _to_list(jp)] | |
| pr = getattr(skel, "joint_prerotations", None) | |
| if pr is not None: | |
| out["prerotations"] = np.asarray(_to_list(pr), float) | |
| print("[sam3d-skel] prerotations shape:", out["prerotations"].shape, | |
| "row0:", out["prerotations"].reshape(out["prerotations"].shape[0], -1)[0].tolist(), flush=True) | |
| off = getattr(skel, "joint_translation_offsets", None) | |
| if off is not None: | |
| out["offsets"] = np.asarray(_to_list(off), float) | |
| print(f"[sam3d-skel] names={None if out['names'] is None else len(out['names'])} " | |
| f"parents={None if out['parents'] is None else len(out['parents'])}", flush=True) | |
| except Exception as e: | |
| print("[sam3d-skel] introspection failed:", repr(e), flush=True) | |
| return out | |
| def run_sam3d_frames(video_path, max_frames=120, log=True): | |
| """Return (global_rots [F,127,3,3], skeleton dict, src_fps). Uses the primary person per frame.""" | |
| import cv2 | |
| _ensure_repo() | |
| est = _estimator() | |
| skel = _get_skeleton(est) | |
| faces = None | |
| try: | |
| faces = np.asarray(_to_list(getattr(est, "faces", None))) | |
| except Exception: | |
| pass | |
| cap = cv2.VideoCapture(video_path) | |
| src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 | |
| rots, mesh, i = [], None, 0 | |
| while i < max_frames: | |
| ok, bgr = cap.read() | |
| if not ok: | |
| break | |
| outs = est.process_one_image(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)) | |
| if log and i == 0 and outs: | |
| _log_structure(outs) | |
| if outs and outs[0].get("pred_global_rots") is not None: | |
| rots.append(np.asarray(_to_list(outs[0]["pred_global_rots"]))) | |
| if mesh is None and faces is not None and outs[0].get("pred_vertices") is not None: | |
| mesh = (np.asarray(_to_list(outs[0]["pred_vertices"])), faces) # frame-0 showcase mesh | |
| i += 1 | |
| cap.release() | |
| g = np.stack(rots) if rots else np.zeros((0, 0, 3, 3)) | |
| return g, skel, src_fps, mesh | |