""" Body3DAgent — optional 3D mesh/joint angle recovery via SAM 3D Body. Input: Pose2DResult, list of athlete masks, list of frames (np.ndarray BGR) Output: Body3DResult(used, joints_3d, confidence) Failure: ALWAYS returns Body3DResult(used=False) when enable_3d=False or checkpoint unavailable — this is a normal success path, not an error. Model: facebook/sam-3d-body-dinov3 (840M params, SAM License, GATED). Gated: YES — access GRANTED June 4, 2026. Params: ~0.84B (DINOv3-H+ variant). API (verified from github.com/facebookresearch/sam-3d-body README, Jun 2026): from notebook.utils import setup_sam_3d_body estimator = setup_sam_3d_body(hf_repo_id="facebook/sam-3d-body-dinov3") outputs = estimator.process_one_image(rgb_image) # single RGB np.ndarray # outputs contains MHR joints, body mesh, etc. """ from __future__ import annotations import numpy as np from formscout.types import Pose2DResult, Body3DResult, IngestResult from formscout import config _NOT_USED = Body3DResult( used=False, joints_3d=[], confidence=0.0, notes="3D disabled or checkpoint unavailable", ) # Subsample frames for 3D inference (expensive per-frame) _MAX_3D_FRAMES = 30 class Body3DAgent: """ Optional 3D body joint estimation via SAM 3D Body (MHR rig). Falls back gracefully when unavailable — returning Body3DResult(used=False) is the expected success path for the 2D-only pipeline. """ def __init__(self, enable_3d: bool | None = None): self._enabled = config.ENABLE_3D if enable_3d is None else enable_3d self._estimator = None if self._enabled: self._estimator = self._try_load() def _try_load(self): """ Attempt to load SAM 3D Body from HuggingFace. Returns the estimator object or None on any failure. """ try: from notebook.utils import setup_sam_3d_body # noqa: F401 estimator = setup_sam_3d_body( hf_repo_id=config.SAM_3D_HF_REPO, ) return estimator except ImportError: return None except Exception: return None def run( self, pose2d: Pose2DResult, masks: list, frames: list | None = None, ) -> Body3DResult: """ Run 3D body estimation on selected keyframes. Args: pose2d: 2D pose results (used for confidence weighting) masks: Per-frame athlete masks from SegmentationAgent frames: Raw BGR frames from IngestResult.frames Returns: Body3DResult with used=True and 3D joints if successful, or Body3DResult(used=False) if disabled/unavailable (normal path). """ if not self._enabled or self._estimator is None: return _NOT_USED if not frames: return Body3DResult( used=False, joints_3d=[], confidence=0.0, notes="3D enabled but no frames provided", ) try: import cv2 # Subsample frames evenly for 3D (it's expensive per-image) n_frames = len(frames) step = max(1, n_frames // _MAX_3D_FRAMES) selected_indices = list(range(0, n_frames, step))[:_MAX_3D_FRAMES] joints_3d_per_frame: list[dict] = [] confidences: list[float] = [] for idx in selected_indices: frame_bgr = frames[idx] # SAM 3D Body expects RGB frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) outputs = self._estimator.process_one_image(frame_rgb) # Extract MHR joint positions from outputs # The model returns joints in the MHR (Momentum Human Rig) format frame_joints = self._extract_joints(outputs, idx) joints_3d_per_frame.append(frame_joints) # Confidence from detection quality conf = self._estimate_confidence(outputs) confidences.append(conf) # Apply light temporal smoothing to reduce jitter joints_3d_smoothed = self._temporal_smooth(joints_3d_per_frame) overall_conf = float(np.mean(confidences)) if confidences else 0.0 return Body3DResult( used=True, joints_3d=joints_3d_smoothed, confidence=overall_conf, notes=f"3D mesh recovery on {len(selected_indices)}/{n_frames} frames", ) except Exception as e: return Body3DResult( used=False, joints_3d=[], confidence=0.0, notes=f"3D inference failed: {e}", ) def _extract_joints(self, outputs: dict, frame_idx: int) -> dict: """ Extract 3D joint positions from SAM 3D Body outputs. Maps MHR rig joints to a standardized dict format. """ joints: dict = {"frame_index": frame_idx} # SAM 3D Body outputs MHR model params including joint positions # The exact key depends on the model output format if hasattr(outputs, "joints_3d"): joint_data = outputs.joints_3d elif isinstance(outputs, dict) and "joints_3d" in outputs: joint_data = outputs["joints_3d"] elif isinstance(outputs, dict) and "pred_joints" in outputs: joint_data = outputs["pred_joints"] else: # Fallback: extract from vertices/body model params joint_data = None if joint_data is not None: if hasattr(joint_data, "cpu"): joint_data = joint_data.cpu().numpy() if isinstance(joint_data, np.ndarray): # Map to named joints (MHR has standard SMPL-like ordering) joint_names = [ "pelvis", "left_hip", "right_hip", "spine1", "left_knee", "right_knee", "spine2", "left_ankle", "right_ankle", "spine3", "left_foot", "right_foot", "neck", "left_collar", "right_collar", "head", "left_shoulder", "right_shoulder", "left_elbow", "right_elbow", "left_wrist", "right_wrist", ] for i, name in enumerate(joint_names): if i < len(joint_data): pos = joint_data[i] joints[name] = { "x": float(pos[0]), "y": float(pos[1]), "z": float(pos[2]), } return joints def _estimate_confidence(self, outputs) -> float: """Estimate confidence from the SAM 3D Body output quality.""" # If outputs have a confidence/score field, use it if isinstance(outputs, dict): if "confidence" in outputs: return float(outputs["confidence"]) if "score" in outputs: return float(outputs["score"]) # Default: assume reasonable confidence if we got outputs at all return 0.75 def _temporal_smooth( self, joints_3d: list[dict], alpha: float = 0.3 ) -> list[dict]: """ Apply exponential moving average smoothing to 3D joint positions to reduce per-frame jitter from single-image prediction. """ if len(joints_3d) <= 1: return joints_3d smoothed = [joints_3d[0]] for i in range(1, len(joints_3d)): prev = smoothed[-1] curr = joints_3d[i] smooth_frame = {"frame_index": curr.get("frame_index", i)} for key in curr: if key == "frame_index": continue if key in prev and isinstance(curr[key], dict) and isinstance(prev[key], dict): smooth_frame[key] = { "x": alpha * curr[key]["x"] + (1 - alpha) * prev[key]["x"], "y": alpha * curr[key]["y"] + (1 - alpha) * prev[key]["y"], "z": alpha * curr[key]["z"] + (1 - alpha) * prev[key]["z"], } else: smooth_frame[key] = curr[key] smoothed.append(smooth_frame) return smoothed