.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_tutorials/sam_3d_body/sam_3d_body.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_tutorials_sam_3d_body_sam_3d_body.py: SAM 3D Body on TensorRT: an end-to-end INT8 deployment tutorial =============================================================== `SAM 3D Body `_ is Meta's single-image full-body human-mesh-recovery model (CVPR 2026). It pairs a **DINOv3 ViT-H+/16 encoder** (840 M parameters — the dominant cost) with a promptable **decoder + MHR parametric head** that regresses an 18439-vertex human mesh. This tutorial deploys the SAM 3D Body model with TensorRT for fast inference on NVIDIA GPUs using `embedl-deploy-tensorrt`. The pipeline consists of: 1. Fetching the HuggingFace model and GitHub code. 2. Quantizing the encoder with the Embedl-Deploy PTQ pipeline. 3. Making the decoder + MHR head ``torch.onnx``-exportable. 4. Building strongly-typed TensorRT engines for both with TensorRT 11. 5. Running a real image end-to-end on the engines, verifying fidelity against eager PyTorch, benchmarking, and rendering the recovered 3D mesh. Results on an NVIDIA L4 (TensorRT 11.1, 512x512 input, batch 1):: encoder INT8+FP16 ~34 ms (FP32 eager PyTorch: ~237 ms -> 7.0x) decoder FP16 ~16 ms total ~50 ms full-pipeline mesh error vs FP32 eager : 2.3 mm mean / 9.7 mm max Requirements (an NVIDIA GPU with ~20 GB VRAM; ~15 GB of disk for artifacts):: pip install torch tensorrt pycuda numpy pillow opencv-python matplotlib pip install imageio onnx huggingface_hub pip install pytorch-lightning yacs scikit-image einops pip install timm roma loguru optree termcolor pip install braceexpand omegaconf onnxscript pip install "embedl-deploy["tensorrt"]" The ``facebook/sam-3d-body-dinov3`` checkpoint is HuggingFace-gated: accept the license on the model page and log in (``hf auth login``) first. This tutorial ships as two files that must sit in the same directory: this script and its export-utilities companion :download:`sam_3d_body_export.py `. Note that the download buttons at the bottom of this page bundle only this script — fetch the companion module through the link above. Download both, then run:: python sam_3d_body.py Every stage caches its artifact in ``./artifacts/`` and is skipped when the artifact already exists, so an interrupted run resumes where it stopped. .. GENERATED FROM PYTHON SOURCE LINES 57-67 Stage 0 — configuration and public sources ------------------------------------------ Everything the pipeline needs is public: * ``facebook/sam-3d-body-dinov3`` (gated): ``model.ckpt`` with all weights, and ``assets/mhr_model.pt``, the TorchScript MHR body rig. * ``facebookresearch/dinov3`` (GitHub): the encoder *architecture* code, loaded through ``torch.hub`` with ``pretrained=False`` — the encoder weights come from the SAM-3D-Body checkpoint, not the DINOv3 release. * ``facebookresearch/sam-3d-body`` (GitHub): the decoder/head code. .. GENERATED FROM PYTHON SOURCE LINES 67-267 .. code-block:: Python import json import statistics import subprocess import sys import time import types import warnings from collections.abc import Callable from pathlib import Path import cv2 # type: ignore[import-not-found] import imageio.v2 as imageio # type: ignore[import-not-found] import matplotlib as mpl # type: ignore[import-not-found] import numpy as np # type: ignore[import-not-found] import numpy.typing as npt # type: ignore[import-not-found] import onnx import pycuda.autoinit # type: ignore[import-not-found] # noqa: F401 import pycuda.driver as cuda # type: ignore[import-not-found] import tensorrt as trt import torch from huggingface_hub import hf_hub_download # type: ignore[import-not-found] from onnx import TensorProto, numpy_helper from PIL import Image # type: ignore[import-not-found] from torch import nn from embedl_deploy import transform from embedl_deploy.backend import set_backend set_backend("tensorrt") from embedl_deploy.quantize import ( CalibrationMethod, ModulesToSkip, QuantConfig, TensorQuantConfig, quantize, ) from embedl_deploy.tensorrt import TENSORRT_PATTERNS mpl.use("Agg") import matplotlib.pyplot as plt # type: ignore[import-not-found] from matplotlib.axes import Axes # type: ignore[import-not-found] from matplotlib.collections import ( # type: ignore[import-not-found] PolyCollection, ) # the SAM-3D-Body code (decoder, MHR head, estimator utilities), pinned to # a known commit so runs are reproducible and the imported code is exactly # the reviewed revision HERE = Path(__file__).resolve().parent REPO_URL = "https://github.com/facebookresearch/sam-3d-body" REPO_COMMIT = "b5c765a0d89d789985e186d396315e7590887b94" DINOV3_COMMIT = "346f38fee679c56a6888f91c51670fae61d364e0" REPO_DIR = HERE / "sam-3d-body" if not REPO_DIR.exists(): print(f"fetching facebookresearch/sam-3d-body @ {REPO_COMMIT[:12]} ...") subprocess.run(["git", "init", "-q", str(REPO_DIR)], check=True) subprocess.run( [ "git", "-C", str(REPO_DIR), "fetch", "-q", "--depth", "1", REPO_URL, REPO_COMMIT, ], check=True, ) subprocess.run( ["git", "-C", str(REPO_DIR), "checkout", "-q", REPO_COMMIT], check=True, ) repo_head = subprocess.run( ["git", "-C", str(REPO_DIR), "rev-parse", "HEAD"], check=True, capture_output=True, text=True, ).stdout.strip() if repo_head != REPO_COMMIT: print( f" warning: sam-3d-body checkout is at {repo_head[:12]}, " f"expected {REPO_COMMIT[:12]}" ) sys.path.insert(0, str(REPO_DIR)) from sam_3d_body import load_sam_3d_body # type: ignore[attr-defined] from sam_3d_body.models.heads import ( # type: ignore[import-not-found] mhr_head as MH, ) from sam_3d_body.models.modules import ( # type: ignore[import-not-found] mhr_utils as MU, ) from sam_3d_body_export import ( _camera_encoder_fwd_nodownsample, _patched_embed_keypoints, _patched_mhr_head_forward, _patched_perspective_projection, _patched_replace_hands_in_pose, _safe_compact_cont_to_model_params_body, _safe_compact_cont_to_model_params_hand, flatten_rig_for_export, register_ts_exporter_symbolics, ) warnings.filterwarnings("ignore") torch.manual_seed(0) ART = HERE / "artifacts" ART.mkdir(exist_ok=True) IMG_SIZE = 512 # SAM-3D-Body production resolution (32x32 tokens) N_CALIB = 16 # PTQ calibration batches SAMPLE_IMAGE_URL = "https://github.com/facebookresearch/sam-3d-body/raw/main/assets/qualitative_comparisons/sample1/input_bbox.png" SAMPLE_IMG = HERE / "sample_input.png" FP32_PT2 = ART / "sam3dbody_backbone_fp32.pt2" REF_FEATURES = ART / "ref_features_fp32.npy" ENC_ONNX = ART / "embedl_sam3dbody_int8.onnx" DEC_ONNX = ART / "sam3dbody_decoder.onnx" FACES_NPY = ART / "sam3dbody_faces.npy" REF_MESH = ART / "ref_mesh_fp32.npz" STATS_PATH = ART / "sam3dbody_trt_stats.json" IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], np.float32).reshape(3, 1, 1) IMAGENET_STD = np.array([0.229, 0.224, 0.225], np.float32).reshape(3, 1, 1) DEC_INPUT_NAMES = [ "features", "rays_ds", "condition_info", "keypoints", "bbox_center", "bbox_scale", "ori_img_size", "cam_int", "affine_trans", "img_size", ] DEC_OUTPUT_NAMES = [ "pred_vertices", "pred_cam_t", "focal_length", "pred_keypoints_3d", "pred_keypoints_2d", "global_rot", "body_pose", "scale", "shape", "hand", "pred_pose_raw", "pred_cam", ] STATS: dict[str, float | int | str] = {} print("=" * 70) print(f"torch : {torch.__version__}") print(f"GPU : {torch.cuda.get_device_name(0)}") # the gated checkpoint (weights + TorchScript MHR rig) print("fetching facebook/sam-3d-body-dinov3 (cached after first run) ...") CKPT_PATH = hf_hub_download("facebook/sam-3d-body-dinov3", "model.ckpt") MHR_PATH = hf_hub_download( "facebook/sam-3d-body-dinov3", "assets/mhr_model.pt" ) # read by load_sam_3d_body hf_hub_download("facebook/sam-3d-body-dinov3", "model_config.yaml") if not SAMPLE_IMG.exists(): print(f"fetching sample image {SAMPLE_IMAGE_URL} ...") imageio.imwrite(SAMPLE_IMG, imageio.imread(SAMPLE_IMAGE_URL)) def benchmark_ms( fn: Callable[[], object], n: int = 50, warmup: int = 10 ) -> tuple[float, float]: """Mean and p95 wall-clock latency of `fn` in milliseconds. :param fn: The callable to time; it must block until its work completes (synchronize the GPU inside). :param n: Number of timed iterations. :param warmup: Number of untimed warmup iterations. :returns: The ``(mean, p95)`` latency in milliseconds. """ for _ in range(warmup): fn() times = [] for _ in range(n): t0 = time.perf_counter() fn() times.append((time.perf_counter() - t0) * 1e3) return statistics.mean(times), sorted(times)[int(0.95 * n)] .. GENERATED FROM PYTHON SOURCE LINES 268-276 Preprocessing — a numpy port of the upstream top-down transforms ---------------------------------------------------------------- SAM 3D Body is a top-down method: it takes a person bounding box, expands it by 1.25x, fixes the aspect ratio, and warps the crop to 512x512. The decoder is additionally conditioned on the *camera*: a ray direction per feature cell (computed from a default full-image intrinsic matrix) and a CLIFF-style bbox-location vector. We reproduce those transforms in plain numpy so the deployed pipeline has no PyTorch dependency at inference time. .. GENERATED FROM PYTHON SOURCE LINES 276-418 .. code-block:: Python def read_rgb(path: Path) -> npt.NDArray[np.uint8]: """Load an image file as an RGB uint8 array.""" img = cv2.imread(str(path)) assert img is not None, f"cannot read image {path}" return np.asarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), dtype=np.uint8) def _fix_aspect_ratio( scale: npt.NDArray[np.float32], ratio: float ) -> npt.NDArray[np.float32]: """Grow the smaller side of `scale` to match the aspect `ratio`.""" w, h = scale return np.array( [w, w / ratio] if w > h * ratio else [h * ratio, h], np.float32 ) def _third( a: npt.NDArray[np.float32], b: npt.NDArray[np.float32] ) -> npt.NDArray[np.float32]: """Return ``b + (a - b)`` rotated 90 degrees anticlockwise.""" d = a - b return b + np.array([-d[1], d[0]], np.float32) def _warp_matrix( center: npt.NDArray[np.float32], scale: npt.NDArray[np.float32], out: int = IMG_SIZE, ) -> npt.NDArray[np.float64]: """Upstream ``get_warp_matrix`` with rot=0 (bbox_utils).""" # numpy broadcast add, not list concatenation src = np.stack([center, center + [0, -0.5 * scale[0]]]) # noqa: RUF005 src = src.astype(np.float32) dst = np.array([[out / 2, out / 2], [out / 2, 0.0]], np.float32) src = np.vstack([src, _third(src[0], src[1])]) dst = np.vstack([dst, _third(dst[0], dst[1])]) # float64, as upstream return np.asarray(cv2.getAffineTransform(src, dst), dtype=np.float64) def preprocess( image_path: Path, bbox: tuple[int, int, int, int] | None = None ) -> tuple[ npt.NDArray[np.float32], dict[str, npt.NDArray[np.float32]], npt.NDArray[np.uint8], ]: """Turn an image and person bbox into encoder and decoder inputs. :param image_path: Path to the input image. :param bbox: Person bounding box ``(x1, y1, x2, y2)``; the full image when omitted. :returns: The normalized ``(1, 3, 512, 512)`` encoder input, the decoder conditioning feed dict (everything except ``features``), and the RGB crop for visualization. """ img = read_rgb(image_path) h, w = img.shape[:2] x1, y1, x2, y2 = bbox or (0, 0, w, h) center = np.array([(x1 + x2) / 2, (y1 + y2) / 2], np.float32) # GetBBoxCenterScale then TopdownAffine, as in the upstream pipeline scale = np.array([x2 - x1, y2 - y1], np.float32) * 1.25 scale = _fix_aspect_ratio(_fix_aspect_ratio(scale, 0.75), 1.0) A = _warp_matrix(center, scale) crop = np.asarray( cv2.warpAffine(img, A, (IMG_SIZE, IMG_SIZE), flags=cv2.INTER_LINEAR), dtype=np.uint8, ) x = ( crop.transpose(2, 0, 1).astype(np.float32) / 255.0 - IMAGENET_MEAN ) / IMAGENET_STD # default full-image intrinsics: focal = image diagonal, principal point # at the image center (upstream's no-calibration fallback) focal = float(np.hypot(w, h)) cam = np.array( [[[focal, 0, w / 2], [0, focal, h / 2], [0, 0, 1]]], np.float32 ) # camera-ray conditioning: the ray (x,y) for every crop pixel, antialias- # downsampled to the 32x32 feature grid (matches the upstream # CameraEncoder's F.interpolate(..., antialias=True)) gx, gy = np.meshgrid( np.arange(IMG_SIZE), np.arange(IMG_SIZE), indexing="xy" ) grid = np.stack([gx, gy], axis=2).astype(np.float32) A32 = A.astype(np.float32) diag, trans = A32[[0, 1], [0, 1]], A32[[0, 1], [2, 2]] grid = grid / diag - trans / diag # crop pixel -> image pixel grid = (grid - cam[0, [0, 1], [2, 2]]) / cam[0, [0, 1], [0, 1]] # -> ray n = IMG_SIZE // 16 rays_ds = np.stack( [ np.asarray( Image.fromarray(grid[:, :, c]).resize( (n, n), Image.Resampling.BILINEAR ) ) for c in range(2) ] )[None] # CLIFF condition, intrinsic-center variant cond = np.array( [ [ (center[0] - cam[0, 0, 2]) / focal, (center[1] - cam[0, 1, 2]) / focal, scale[0] / focal, ] ], np.float32, ) feeds = { "rays_ds": rays_ds.astype(np.float32), "condition_info": cond, "keypoints": np.array([[[0.0, 0.0, -2.0]]], np.float32), # dummy "bbox_center": center.reshape(1, 1, 2), "bbox_scale": scale.reshape(1, 1, 2), "ori_img_size": np.array([[[w, h]]], np.float32), "cam_int": cam, "affine_trans": A32.reshape(1, 1, 2, 3), "img_size": np.full((1, 1, 2), IMG_SIZE, np.float32), } return x[None], feeds, crop print("\n[stage 0] preprocessing the sample image ...") enc_input, dec_feeds, crop_rgb = preprocess(SAMPLE_IMG) print( f" encoder input {enc_input.shape}, decoder conditioning " f"{sorted(dec_feeds)}" ) .. GENERATED FROM PYTHON SOURCE LINES 419-431 Stage 1 — the encoder, with the real checkpoint weights ------------------------------------------------------- The checkpoint stores the whole model; the encoder lives under the ``backbone.encoder.`` prefix. We wrap it: plain tensors in, plain tensors out, which is what ``torch.export`` and the Embedl PTQ flow expect. Why deploy the encoder *separately* from the decoder? The full ``SAM3DBody.forward`` mixes data-dependent control flow (body/hand batch selection), so it doesn't fit a single static graph. The encoder is a clean image-in/features-out function and carries most of the compute, so it is the natural quantization target; the decoder is exported as a second, FP16-only graph in stage 3. .. GENERATED FROM PYTHON SOURCE LINES 431-516 .. code-block:: Python class BackboneWrapper(nn.Module): """Tensor-in/tensor-out wrapper around the DINOv3 encoder.""" def __init__(self, encoder: nn.Module) -> None: """Wrap the DINOv3 `encoder`.""" super().__init__() self.encoder = encoder def forward(self, image: torch.Tensor) -> torch.Tensor: """Return the final 32x32 feature map for `image`.""" feats = self.encoder.get_intermediate_layers( # type: ignore[operator] image, n=1, reshape=True, norm=True ) return feats[-1] # type: ignore[no-any-return] # (B, 1280, 32, 32) def load_encoder_with_real_weights(device: str) -> nn.Module: """Build the DINOv3 architecture and load the SAM-3D-Body weights. :param device: Device to place the encoder on. :returns: The encoder in eval mode with the checkpoint weights loaded. """ encoder = torch.hub.load( # type: ignore[no-untyped-call] f"facebookresearch/dinov3:{DINOV3_COMMIT}", "dinov3_vith16plus", source="github", pretrained=False, trust_repo=True, skip_validation=True, # ref is a pinned commit, not a branch ).eval() ckpt = torch.load(CKPT_PATH, map_location="cpu", weights_only=False) state = ckpt.get("state_dict", ckpt) backbone_state = { k.removeprefix("backbone.encoder."): v for k, v in state.items() if k.startswith("backbone.encoder.") } missing, unexpected = encoder.load_state_dict(backbone_state, strict=False) # the checkpoint omits only the mask_token, which inference never uses assert set(missing) <= {"mask_token"}, f"missing encoder keys: {missing}" assert not unexpected, f"checkpoint keys not in encoder: {unexpected}" return encoder.eval().to(device) # type: ignore[no-any-return] if not (FP32_PT2.exists() and REF_FEATURES.exists()): print("\n[stage 1] loading the DINOv3 ViT-H+/16 encoder (real weights)...") model = BackboneWrapper(load_encoder_with_real_weights("cuda")) n_params = sum(p.numel() for p in model.parameters()) print( f" {n_params / 1e6:.1f}M params, {n_params * 4 / 1e9:.2f} GB FP32" ) x = torch.from_numpy(enc_input).cuda() with torch.no_grad(): ref = model(x) np.save(REF_FEATURES, ref.float().cpu().numpy()) print(f" FP32 reference features {tuple(ref.shape)} saved") # eager FP32 latency -- the baseline the TRT engine is measured against def _run_eager() -> None: with torch.no_grad(): model(x) torch.cuda.synchronize() STATS["encoder_fp32_eager_ms"], _ = benchmark_ms( _run_eager, n=20, warmup=5 ) print(f" eager FP32 latency: {STATS['encoder_fp32_eager_ms']:.1f} ms") print(" torch.export.export(encoder) ...") dummy = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, device="cuda") with torch.no_grad(): exported = torch.export.export(model, (dummy,), strict=False) torch.export.save(exported, str(FP32_PT2)) print(f" saved {FP32_PT2.name}") del model, exported, ref torch.cuda.empty_cache() else: print("\n[stage 1] reusing cached FP32 export + reference features") .. GENERATED FROM PYTHON SOURCE LINES 517-548 Stage 2 — INT8 post-training quantization with Embedl-Deploy ------------------------------------------------------------ The standard Embedl-Deploy recipe — ``transform`` fuses the exported graph into TensorRT-friendly patterns, ``quantize`` inserts Q/DQ nodes and calibrates them — plus three customizations that matter for this transformer: * **Skip the patch-embed conv.** The 3-channel input projection gains almost nothing from INT8 and loses accuracy disproportionately. * **Don't SmoothQuant the LayerNorms.** Smoothing folds scales into LayerNorm weights in-place; on this backbone that hurts more than it helps, so LayerNorm modules are excluded via ``ModulesToSkip(smooth=...)``. * **Calibrate on the real crop.** MINMAX with 16 lightly-jittered copies of the actual person crop gives ranges that match deployment. .. note:: To keep the tutorial self-contained, PTQ calibrates on jittered copies of a single image. In a realistic deployment you would calibrate on a few hundred representative person crops from your target domain — for example a sample of the datasets SAM 3D Body itself is trained and evaluated on, such as COCO, MPII, AI Challenger, 3DPW, EgoHumans, EgoExo4D, Harmony4D, and SA-1B (see ``data/scripts`` in the upstream ``sam-3d-body`` repository for download helpers). Note that the *final* INT8 coverage is decided by the backend's precision propagation inside ``quantize``: sites whose TensorRT neighbours cannot run INT8 (e.g., around scaled-dot-product attention) are settled back to FP16, so the exported graph carries Q/DQ only where TensorRT will actually use it. The quantized graph is re-exported with ``torch.export`` and saved to ONNX with external data (the model is >2 GB, above protobuf's limit). .. GENERATED FROM PYTHON SOURCE LINES 548-661 .. code-block:: Python def export_onnx_external( module: nn.Module, args: tuple[torch.Tensor, ...], onnx_path: Path ) -> None: """Dynamo-export to ONNX with all tensors in one external .data file. :param module: The module to export. :param args: Example positional inputs for the export. :param onnx_path: Destination path; the tensor data lands next to it in ``.data``. """ tmp = onnx_path.with_suffix(".tmp.onnx") with torch.no_grad(): torch.onnx.export( module, args, str(tmp), input_names=["image"], output_names=["features"], do_constant_folding=True, dynamo=True, ) m = onnx.load(str(tmp), load_external_data=True) onnx.save_model( m, str(onnx_path), save_as_external_data=True, all_tensors_to_one_file=True, location=onnx_path.name + ".data", size_threshold=1024, ) tmp.unlink(missing_ok=True) for p in onnx_path.parent.glob(tmp.name + "*"): p.unlink(missing_ok=True) if not ENC_ONNX.exists(): print("\n[stage 2] Embedl-Deploy transform + INT8 PTQ ...") dummy = torch.randn(1, 3, IMG_SIZE, IMG_SIZE, device="cuda") graph = torch.export.load(str(FP32_PT2)).module() fused = transform(graph, (dummy,), TENSORRT_PATTERNS).model.eval() # the only 3-channel conv in the graph is the 16x16 patch embed patch_embed = next( m for m in fused.modules() if isinstance(m, nn.Conv2d) and m.in_channels == 3 ) quant_cfg = QuantConfig( activation=TensorQuantConfig( # INT8, per-tensor activations symmetric=True, per_channel=False, calibration_method=CalibrationMethod.MINMAX, ), weight=TensorQuantConfig( # INT8, per-channel weights symmetric=True, per_channel=True, calibration_method=CalibrationMethod.MINMAX, ), skip=ModulesToSkip( stub={patch_embed}, weight={patch_embed}, smooth={nn.LayerNorm}, ), ) real = torch.from_numpy(enc_input).cuda() calib = [real + 0.02 * torch.randn_like(real) for _ in range(N_CALIB)] def calibrate(m: nn.Module) -> None: """Run the calibration batches through the model.""" m.eval() with torch.no_grad(): for c in calib: m(c) print(f" calibrating on {N_CALIB} jittered copies of the real crop ...") qmodel = quantize( fused, (dummy,), config=quant_cfg, forward_loop=calibrate, freeze_weights=True, ) # fidelity of the QDQ graph itself (before TensorRT enters the picture) with torch.no_grad(): q_feats = qmodel(real).float().cpu().numpy() ref = np.load(REF_FEATURES) rel = np.abs(q_feats - ref).mean() / (np.abs(ref).mean() + 1e-9) STATS["encoder_qdq_rel_diff"] = float(rel) print(f" QDQ-vs-FP32 feature rel diff: {rel * 100:.3f}%") print(" exporting INT8 ONNX (external data) ...") with torch.no_grad(): q_exported = torch.export.export(qmodel, (dummy,), strict=False) export_onnx_external(q_exported.module(), (dummy,), ENC_ONNX) sz = ( ENC_ONNX.stat().st_size + (ART / (ENC_ONNX.name + ".data")).stat().st_size ) print(f" saved {ENC_ONNX.name} (+ .data, {sz / 1e9:.2f} GB)") del graph, fused, qmodel, q_exported torch.cuda.empty_cache() else: print("\n[stage 2] reusing cached INT8 encoder ONNX") .. GENERATED FROM PYTHON SOURCE LINES 662-688 Stage 3 — making the decoder + MHR head exportable -------------------------------------------------- The decoder (no-mask prompt add, ray-conditioned camera encoder, 6 promptable decoder layers with per-layer MHR head + camera projection) is wrapped with pure tensor I/O, but cannot be exported as-is: 1. **The MHR rig is TorchScript.** ``torch.jit.load``-ed modules can't be traced through. We run the rig once under ``make_fx(functionalize(...))`` to record the dispatched aten ops into a flat FX graph, then clean up the artifacts of that flattening (constant sparse tensors -> dense, in-place copies -> broadcasts, index_copy -> scatter with baked indices). 2. **Advanced-index assignment** (``x[:, idx] = v``, boolean-mask writes) exports through a broken ``index_put`` lowering. Every such site is replaced by a numerically-identical gather/scatter/where formulation. 3. **roma's rotmat->euler** handles gimbal lock with data-dependent ``nonzero()`` indexing, which bakes trace-time element counts into the graph. A branchless torch.where version replaces it. 4. **Antialiased ray downsampling** inside the CameraEncoder — TensorRT's Resize has no antialias, so rays are downsampled on the host (stage 0) and passed in at feature resolution. The machinery that fixes all four — FX graph surgery, ONNX-safe method patches, and custom exporter symbolics — lives in the companion module ``sam_3d_body_export.py``. Each patch is numerically identical to the code it replaces (the rig flattening self-checks on five probe input sets before it is accepted), so here we only *apply* it and export. .. GENERATED FROM PYTHON SOURCE LINES 688-845 .. code-block:: Python class BodyDecoderWrapper(nn.Module): """No-mask prompt add + camera encoder + 6-layer promptable decoder (with per-layer MHR head, camera projection, and keypoint-token updates) -> final pose/mesh outputs. Batch size 1, one dummy keypoint prompt (label -2), no mask — exactly the body-only inference configuration. """ def __init__(self, model: nn.Module) -> None: """Wrap the full SAM3DBody `model`.""" super().__init__() self.model = model def forward( # noqa: PLR0913 -- one argument per ONNX input self, features: torch.Tensor, rays_ds: torch.Tensor, condition_info: torch.Tensor, keypoints: torch.Tensor, bbox_center: torch.Tensor, bbox_scale: torch.Tensor, ori_img_size: torch.Tensor, cam_int: torch.Tensor, affine_trans: torch.Tensor, img_size: torch.Tensor, ) -> tuple[torch.Tensor, ...]: """Run the decoder on one person crop. Input shapes:: features: (1, 1280, 32, 32) backbone output rays_ds: (1, 2, 32, 32) downsampled camera rays condition_info: (1, 3) CLIFF condition keypoints: (1, 1, 3) prompt (dummy: [0, 0, -2]) bbox_center: (1, 1, 2) bbox_scale: (1, 1, 2) ori_img_size: (1, 1, 2) cam_int: (1, 3, 3) affine_trans: (1, 1, 2, 3) img_size: (1, 1, 2) """ m = self.model m._batch_size, m._max_num_person = 1, 1 # type: ignore[assignment] # noqa: SLF001 m.body_batch_idx, m.hand_batch_idx = [0], [] # type: ignore[assignment] # use_mask=False path of _get_mask_prompt reduces to a constant add no_mask = m.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1) # type: ignore[operator, union-attr] features = features + no_mask batch = { # only img.shape[1] (num_person) is read "img": features.new_zeros(1, 1, 1), "ray_cond": rays_ds, "bbox_center": bbox_center, "bbox_scale": bbox_scale, "ori_img_size": ori_img_size, "cam_int": cam_int, "affine_trans": affine_trans, "img_size": img_size, } _tokens, pose_outputs = m.forward_decoder( # type: ignore[operator] features, init_estimate=None, keypoints=keypoints, prev_estimate=None, condition_info=condition_info, batch=batch, ) po = pose_outputs[-1] return tuple(po[k] for k in DEC_OUTPUT_NAMES) def apply_export_patches(model: nn.Module) -> None: """Swap every export-hostile site for an ONNX-safe equivalent. Every patch is numerically identical to the code it replaces for the body-only inference configuration the tutorial exports. :param model: The ``SAM3DBody`` model, patched in place. """ # rebind in both namespaces: mhr_head imports these by name, so patching # the mhr_utils module attribute alone would not reach its call sites for mod in (MU, MH): mod.compact_cont_to_model_params_body = ( _safe_compact_cont_to_model_params_body ) mod.compact_cont_to_model_params_hand = ( _safe_compact_cont_to_model_params_hand ) model.ray_cond_emb.forward = types.MethodType( # type: ignore[union-attr] _camera_encoder_fwd_nodownsample, model.ray_cond_emb ) model.head_pose.replace_hands_in_pose = types.MethodType( # type: ignore[union-attr] _patched_replace_hands_in_pose, model.head_pose ) model.head_pose.forward = types.MethodType( # type: ignore[union-attr] _patched_mhr_head_forward, model.head_pose ) model.head_camera.perspective_projection = types.MethodType( # type: ignore[union-attr] _patched_perspective_projection, model.head_camera ) model.prompt_encoder._embed_keypoints = types.MethodType( # type: ignore[union-attr] # noqa: SLF001 _patched_embed_keypoints, model.prompt_encoder ) if not (DEC_ONNX.exists() and FACES_NPY.exists() and REF_MESH.exists()): print("\n[stage 3] exporting the decoder + MHR head to ONNX ...") # everything on CPU: the flattened rig bakes its trace device into the # graph, and the ONNX constant folder needs single-device constants body_model, _cfg = load_sam_3d_body( CKPT_PATH, device="cpu", mhr_path=MHR_PATH ) np.save(FACES_NPY, body_model.head_pose.faces.cpu().numpy()) apply_export_patches(body_model) wrapper = BodyDecoderWrapper(body_model).eval() ref_features = np.load(REF_FEATURES) dec_args = tuple( torch.from_numpy({"features": ref_features, **dec_feeds}[k]) for k in DEC_INPUT_NAMES ) print(" flattening the TorchScript MHR rig to FX ...") flatten_rig_for_export(wrapper, body_model.head_pose, MHR_PATH, dec_args) # FP32 eager reference mesh for the real sample -- the ground truth that # the TensorRT pipeline is compared against in stage 5 with torch.no_grad(): ref_out = dict(zip(DEC_OUTPUT_NAMES, wrapper(*dec_args), strict=True)) np.savez( REF_MESH, vertices=ref_out["pred_vertices"][0].numpy(), cam_t=ref_out["pred_cam_t"][0].numpy(), ) print( f" FP32 reference mesh: {ref_out['pred_vertices'].shape[1]} " f"vertices" ) register_ts_exporter_symbolics() print(" exporting ONNX (TorchScript tracer, CPU) ...") with torch.no_grad(): torch.onnx.export( wrapper, dec_args, str(DEC_ONNX), input_names=DEC_INPUT_NAMES, output_names=DEC_OUTPUT_NAMES, opset_version=17, dynamo=False, do_constant_folding=True, ) print( f" saved {DEC_ONNX.name} ({DEC_ONNX.stat().st_size / 1e6:.0f} MB)" ) del body_model, wrapper else: print("\n[stage 3] reusing cached decoder ONNX + reference mesh") .. GENERATED FROM PYTHON SOURCE LINES 846-859 Stage 4 — TensorRT engines -------------------------- TensorRT 11 builds **strongly-typed** networks: every layer keeps the precision declared in the ONNX file, and the old FP16/INT8 builder flags are gone. So we set precisions in the ONNX itself: * FP32 tensors, Cast targets, and inline constants are rewritten to FP16 (clipped to the FP16 range). The encoder's Q/DQ INT8 structure is untouched — those layers run INT8, the rest FP16. * The decoder's float64 mesh-rig section (FK chains accumulate error) is demoted to FP32, not FP16, to keep the skinning math accurate. The cast and the engine build each happen once and are cached on disk. .. GENERATED FROM PYTHON SOURCE LINES 859-963 .. code-block:: Python TRT_LOG = trt.Logger(trt.Logger.WARNING) ENC_ENGINE = ART / "embedl_sam3dbody_int8.engine" DEC_ENGINE = ART / "sam3dbody_decoder.engine" ELEM_REMAP = { TensorProto.FLOAT: TensorProto.FLOAT16, TensorProto.DOUBLE: TensorProto.FLOAT, } def _retype_tensor(t: onnx.TensorProto) -> None: """Rewrite a tensor fp32 -> clipped fp16 or fp64 -> fp32, in place.""" if t.data_type == TensorProto.FLOAT: arr = np.clip(numpy_helper.to_array(t), -65504, 65504).astype( np.float16 ) t.CopyFrom(numpy_helper.from_array(arr, t.name)) elif t.data_type == TensorProto.DOUBLE: arr = numpy_helper.to_array(t).astype(np.float32) t.CopyFrom(numpy_helper.from_array(arr, t.name)) def cast_onnx_to_fp16(onnx_path: Path) -> Path: """Cast an fp32 ONNX file to fp16 for a strongly-typed TensorRT build. fp64 becomes fp32 and fp32 becomes fp16 across initializers, value infos, Cast targets, and inline tensor attributes; the Q/DQ INT8 structure is untouched. :param onnx_path: The fp32 ONNX file to cast. :returns: The path of the cast copy (cached across runs). """ fp16_path = onnx_path.with_name(onnx_path.stem + "_fp16.onnx") if fp16_path.exists(): return fp16_path print(f" casting {onnx_path.name} to fp16 (one-time) ...") model = onnx.load(str(onnx_path)) g = model.graph for init in g.initializer: _retype_tensor(init) for vi in list(g.input) + list(g.output) + list(g.value_info): t = vi.type.tensor_type t.elem_type = ELEM_REMAP.get(t.elem_type, t.elem_type) for node in g.node: for a in node.attribute: if node.op_type == "Cast" and a.name == "to": a.i = ELEM_REMAP.get(a.i, a.i) elif a.type == onnx.AttributeProto.TENSOR: _retype_tensor(a.t) onnx.save( model, str(fp16_path), save_as_external_data=True, location=fp16_path.name + ".data", all_tensors_to_one_file=True, ) return fp16_path def build_engine(onnx_path: Path, engine_path: Path) -> trt.ICudaEngine: """Build (or load cached) a strongly-typed fp16 engine from fp32 ONNX. :param onnx_path: The fp32 ONNX file (cast to fp16 on first use). :param engine_path: Where the serialized engine is cached. :returns: The deserialized TensorRT engine. """ trt.init_libnvinfer_plugins(TRT_LOG, "") if engine_path.exists(): return trt.Runtime(TRT_LOG).deserialize_cuda_engine( engine_path.read_bytes() ) onnx_path = cast_onnx_to_fp16(onnx_path) print(f" building {engine_path.name} from {onnx_path.name} ...") builder = trt.Builder(TRT_LOG) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) ) parser = trt.OnnxParser(network, TRT_LOG) if not parser.parse_from_file(str(onnx_path)): raise RuntimeError( "\n".join( str(parser.get_error(i)) for i in range(parser.num_errors) ) ) cfg = builder.create_builder_config() cfg.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * 1024**3) t0 = time.perf_counter() plan = builder.build_serialized_network(network, cfg) if plan is None: raise RuntimeError("engine build failed") engine_path.write_bytes(bytes(plan)) print( f" {engine_path.name}: built in {time.perf_counter() - t0:.0f}s, " f"{engine_path.stat().st_size / 1e6:.0f} MB" ) return trt.Runtime(TRT_LOG).deserialize_cuda_engine(bytes(plan)) .. GENERATED FROM PYTHON SOURCE LINES 964-970 Stage 5 — the TensorRT demo: real image in, 3D mesh out ------------------------------------------------------- A small pycuda runner drives both engines: pinned host buffers, one device buffer per I/O tensor, `execute_async_v3` on a private stream. Latency is measured on engine execution only (inputs already on device), which is the number that matters when the pipeline runs inside a larger system. .. GENERATED FROM PYTHON SOURCE LINES 970-1077 .. code-block:: Python class Engine: """Single-context TensorRT runner with engine-only latency measurement.""" def __init__(self, onnx_path: Path, engine_path: Path) -> None: """Build or load the engine and allocate its I/O buffers.""" self.engine = build_engine(onnx_path, engine_path) self.ctx = self.engine.create_execution_context() self.stream = cuda.Stream() self.host: dict[str, npt.NDArray[np.generic]] = {} self.dev: dict[str, cuda.DeviceAllocation] = {} self.shapes: dict[str, tuple[int, ...]] = {} self.inputs: list[str] = [] self.outputs: list[str] = [] for i in range(self.engine.num_io_tensors): nm = self.engine.get_tensor_name(i) shape = tuple(self.engine.get_tensor_shape(nm)) dt = trt.nptype(self.engine.get_tensor_dtype(nm)) self.host[nm] = cuda.pagelocked_empty(int(np.prod(shape)) or 1, dt) self.dev[nm] = cuda.mem_alloc(max(self.host[nm].nbytes, 8)) self.shapes[nm] = shape self.ctx.set_tensor_address(nm, int(self.dev[nm])) is_in = self.engine.get_tensor_mode(nm) == trt.TensorIOMode.INPUT (self.inputs if is_in else self.outputs).append(nm) def __call__( self, feeds: dict[str, npt.NDArray[np.float32]] ) -> dict[str, npt.NDArray[np.float32]]: """Run one inference on `feeds` and return the fp32 outputs.""" for nm in self.inputs: np.copyto( self.host[nm], feeds[nm].ravel().astype(self.host[nm].dtype) ) cuda.memcpy_htod_async(self.dev[nm], self.host[nm], self.stream) self.ctx.execute_async_v3(self.stream.handle) for nm in self.outputs: cuda.memcpy_dtoh_async(self.host[nm], self.dev[nm], self.stream) self.stream.synchronize() return { nm: self.host[nm] .reshape(self.shapes[nm]) .astype(np.float32) .copy() for nm in self.outputs } def benchmark(self) -> tuple[float, float]: """Mean and p95 engine-only latency in milliseconds.""" def _run() -> None: self.ctx.execute_async_v3(self.stream.handle) self.stream.synchronize() return benchmark_ms(_run) print("\n[stage 4] TensorRT engines ...") encoder = Engine(ENC_ONNX, ENC_ENGINE) decoder = Engine(DEC_ONNX, DEC_ENGINE) print("\n[stage 5] running the sample end to end on TensorRT ...") features = encoder({"image": enc_input})["features"] out = decoder({**dec_feeds, "features": features}) V = out["pred_vertices"][0] cam_t = out["pred_cam_t"][0] print(f" recovered mesh: {V.shape[0]} vertices") # fidelity vs the eager FP32 reference ref_feats = np.load(REF_FEATURES) feat_rel = np.abs(features - ref_feats).mean() / ( np.abs(ref_feats).mean() + 1e-9 ) ref_mesh = np.load(REF_MESH) mesh_d = np.abs(V - ref_mesh["vertices"]) STATS.update( { "encoder_trt_rel_diff": float(feat_rel), "mesh_max_mm": float(mesh_d.max() * 1000), "mesh_mean_mm": float(mesh_d.mean() * 1000), } ) print( f" encoder features INT8-TRT vs FP32 eager: " f"rel diff {feat_rel * 100:.3f}%" ) print( f" full-pipeline mesh vs FP32 eager: " f"max {mesh_d.max() * 1000:.2f} mm mean {mesh_d.mean() * 1000:.3f} mm" ) enc_ms, enc_p95 = encoder.benchmark() dec_ms, dec_p95 = decoder.benchmark() STATS.update( { "encoder_trt_ms": enc_ms, "encoder_trt_p95_ms": enc_p95, "decoder_trt_ms": dec_ms, "decoder_trt_p95_ms": dec_p95, "tensorrt": trt.__version__, "torch": torch.__version__, "gpu": torch.cuda.get_device_name(0), "input_size": IMG_SIZE, } ) .. GENERATED FROM PYTHON SOURCE LINES 1078-1082 Rendering — mesh overlay + turntable ------------------------------------ The mesh is projected back onto the crop with the same intrinsics + affine used in preprocessing, flat-shaded, and painter-sorted (no GL dependency). .. GENERATED FROM PYTHON SOURCE LINES 1082-1234 .. code-block:: Python LIGHT = np.array([0.3, 0.5, 1.0]) LIGHT /= np.linalg.norm(LIGHT) SKIN = np.array([0.80, 0.78, 0.72]) def _shade( v: npt.NDArray[np.float32], f: npt.NDArray[np.int64] ) -> npt.NDArray[np.float64]: """Per-face flat-shading colors from a fixed directional light.""" n = np.cross(v[f][:, 1] - v[f][:, 0], v[f][:, 2] - v[f][:, 0]) n /= np.linalg.norm(n, axis=1, keepdims=True) + 1e-9 lam = np.clip(np.abs(n @ LIGHT), 0, 1)[:, None] return np.clip(0.25 + 0.75 * lam * SKIN, 0, 1) def _view( ax: Axes, verts: npt.NDArray[np.float32], faces: npt.NDArray[np.int64], deg: float, title: str, ) -> None: """Draw the mesh into `ax`, rotated `deg` degrees about the y axis.""" Vc = verts - verts.mean(0) th = np.radians(deg) R = np.array( [[np.cos(th), 0, np.sin(th)], [0, 1, 0], [-np.sin(th), 0, np.cos(th)]] ) Vr = Vc @ R.T p = Vr[:, :2] * [1, -1] o = np.argsort(Vr[faces].mean(1)[:, 2]) ax.add_collection( PolyCollection( p[faces][o], facecolors=_shade(Vr, faces)[o], edgecolors="none" ) ) ax.set_xlim(p[:, 0].min(), p[:, 0].max()) ax.set_ylim(p[:, 1].min(), p[:, 1].max()) ax.set_aspect("equal") ax.axis("off") ax.set_title(title, fontsize=11) def render_overlay( image_path: Path, verts: npt.NDArray[np.float32], faces: npt.NDArray[np.int64], cam_t: npt.NDArray[np.float32], crop: npt.NDArray[np.uint8], cam_int: npt.NDArray[np.float32], affine: npt.NDArray[np.float32], out: Path, ) -> None: """Save a 4-panel figure: input, crop overlay, and two mesh views. :param image_path: Path of the original input image. :param verts: Mesh vertices, shape ``(N, 3)``. :param faces: Mesh faces, shape ``(M, 3)``. :param cam_t: Predicted camera translation, shape ``(3,)``. :param crop: The 512x512 RGB model crop from preprocessing. :param cam_int: Camera intrinsics used in preprocessing, shape ``(1, 3, 3)``. :param affine: The full-image-to-crop affine from preprocessing. :param out: Destination PNG path. """ img = read_rgb(image_path) fig, ax = plt.subplots(1, 4, figsize=(15, 6)) fig.patch.set_facecolor("white") ax[0].imshow(img) ax[0].axis("off") ax[0].set_title("Input", fontsize=11) # project with the camera intrinsics, then map full-image pixels into the # model crop with the same affine used for preprocessing Vc = verts + cam_t z = np.clip(Vc[:, 2], 1e-3, None) f, cx, cy = cam_int[0, 0, 0], cam_int[0, 0, 2], cam_int[0, 1, 2] pf = np.stack( [f * Vc[:, 0] / z + cx, f * Vc[:, 1] / z + cy, np.ones_like(z)], 1 ) A = affine.reshape(2, 3).astype(np.float64) p = pf @ A.T o = np.argsort(-Vc[faces].mean(1)[:, 2]) ax[1].imshow(crop) ax[1].add_collection( PolyCollection( p[faces][o], facecolors=_shade(Vc, faces)[o], edgecolors="none", alpha=0.8, ) ) ax[1].set_xlim(0, IMG_SIZE) ax[1].set_ylim(IMG_SIZE, 0) ax[1].axis("off") ax[1].set_title("Mesh overlay", fontsize=11) _view(ax[2], verts, faces, 20, "¾ view") _view(ax[3], verts, faces, 90, "side view") plt.tight_layout() plt.savefig(out, dpi=160, bbox_inches="tight") plt.close() print(f" wrote {out}") def render_turntable( verts: npt.NDArray[np.float32], faces: npt.NDArray[np.int64], out: Path, ) -> None: """Save a rotating-mesh GIF (one frame every 15 degrees). :param verts: Mesh vertices, shape ``(N, 3)``. :param faces: Mesh faces, shape ``(M, 3)``. :param out: Destination GIF path. """ frames = [] for a in range(0, 360, 15): fig, ax = plt.subplots(figsize=(4, 6)) fig.patch.set_facecolor("white") _view(ax, verts, faces, a, "") fig.canvas.draw() rgba = fig.canvas.buffer_rgba() frames.append(np.asarray(rgba)[..., :3].copy()) plt.close() imageio.mimsave(out, frames, duration=0.1, loop=0) print(f" wrote {out}") faces = np.load(FACES_NPY) render_overlay( SAMPLE_IMG, V, faces, cam_t, crop_rgb, dec_feeds["cam_int"], dec_feeds["affine_trans"], ART / "mesh_demo_trt.png", ) render_turntable(V, faces, ART / "mesh_demo_trt_spin.gif") .. GENERATED FROM PYTHON SOURCE LINES 1235-1237 Summary ------- .. GENERATED FROM PYTHON SOURCE LINES 1237-1275 .. code-block:: Python STATS_PATH.write_text(json.dumps(STATS, indent=2)) eager = STATS.get("encoder_fp32_eager_ms") eager_ms = eager if isinstance(eager, float) else None print("\n" + "=" * 72) print( f"SAM 3D Body on TensorRT {trt.__version__} -- " f"{torch.cuda.get_device_name(0)}, {IMG_SIZE}x{IMG_SIZE}" ) print("=" * 72) print(f"{'Stage':<28}{'Mean':>10}{'p95':>10}{'Engine':>12}") print("-" * 72) if eager_ms: print( f"{'encoder FP32 (eager torch)':<28}{eager_ms:>7.2f} ms" f"{'--':>10}{'--':>12}" ) for name, ms, p95, engine_path in ( ("encoder INT8+FP16 (TRT)", enc_ms, enc_p95, ENC_ENGINE), ("decoder FP16 (TRT)", dec_ms, dec_p95, DEC_ENGINE), ): print( f"{name:<28}{ms:>7.2f} ms{p95:>7.2f} ms" f"{engine_path.stat().st_size / 2**20:>8.0f} MiB" ) print(f"{'total (TRT)':<28}{enc_ms + dec_ms:>7.2f} ms") if eager_ms: print(f"\nencoder speedup vs eager FP32: {eager_ms / enc_ms:.2f}x") print( f"encoder features rel diff (INT8 engine vs FP32 eager): " f"{STATS['encoder_trt_rel_diff'] * 100:.3f}%" ) print( f"mesh error vs FP32 eager pipeline: mean " f"{STATS['mesh_mean_mm']:.3f} mm, max {STATS['mesh_max_mm']:.2f} mm" ) print("=" * 72) print(f"\nresults written to {STATS_PATH}") .. _sphx_glr_download_auto_tutorials_sam_3d_body_sam_3d_body.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: sam_3d_body.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: sam_3d_body.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: sam_3d_body.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_