Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/004_sugar_box.glb filter=lfs diff=lfs merge=lfs -text
+examples/006_mustard_bottle.glb filter=lfs diff=lfs merge=lfs -text
+examples/019_pitcher_base.glb filter=lfs diff=lfs merge=lfs -text
+examples/035_power_drill.glb filter=lfs diff=lfs merge=lfs -text
+examples/mustard_bottle_rgb.png filter=lfs diff=lfs merge=lfs -text
+examples/pitcher_base_rgb.png filter=lfs diff=lfs merge=lfs -text
+examples/power_drill_rgb.png filter=lfs diff=lfs merge=lfs -text
+examples/sugar_box_rgb.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,45 @@
 ---
-title: Foundationpose 6dof
-emoji: 👁
+title: FoundationPose 6-DoF
+emoji: 📦
 colorFrom: green
-colorTo: red
+colorTo: blue
 sdk: gradio
 sdk_version: 6.27.0
-python_version: '3.12'
 app_file: app.py
+short_description: 6-DoF pose of novel objects from RGB-D plus a CAD model
+python_version: "3.12"
+startup_duration_timeout: 45m
 pinned: false
+license: other
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# FoundationPose — 6-DoF pose estimation of novel objects
+
+[FoundationPose](https://arxiv.org/abs/2312.08344) (NVIDIA, CVPR 2024) estimates the
+full 6-DoF pose of an object it has never seen during training, given
+
+* an **RGB image**,
+* the **aligned depth map**,
+* the object's **CAD model**,
+* the **camera intrinsics** (`fx, fy, cx, cy`), and
+* a **2D bounding box** around the object.
+
+The demo reproduces the model-based *registration* path of the reference
+implementation: 252 pose hypotheses are sampled on an icosphere (42 viewpoints ×
+6 in-plane rotations), the translation is initialised from the box centre and the
+median depth inside the box, every hypothesis is refined for 5 passes by the
+refiner network, and the scorer network ranks them.
+
+The published checkpoints ([`nvidia/foundationpose`](https://huggingface.co/nvidia/foundationpose))
+are ONNX; they are converted to PyTorch with
+[onnx2torch](https://github.com/ENOT-AutoDL/onnx2torch) at startup, so the whole
+pipeline runs natively on ZeroGPU. Rendering uses
+[nvdiffrast](https://github.com/NVlabs/nvdiffrast).
+
+## Credits and licences
+
+* Weights: NVIDIA Open Model License — <https://huggingface.co/nvidia/foundationpose>
+* Reference code: [NVlabs/FoundationPose](https://github.com/NVlabs/FoundationPose)
+* Example scenes and CAD models: [YCB-Video / BOP](https://huggingface.co/datasets/bop-benchmark/ycbv),
+  MIT licence, © 2017 University of Washington Robotics and State Estimation Lab —
+  one of FoundationPose's own evaluation datasets.

app.py

@@ -0,0 +1,1023 @@
+"""FoundationPose — 6-DoF pose estimation of novel objects (NVIDIA, CVPR 2024).
+
+Model-based registration path: given an RGB-D frame, the object's CAD model,
+camera intrinsics and a 2D box, sample 252 pose hypotheses on an icosphere,
+refine each with the refiner network, then rank them with the scorer network.
+
+Weights: https://huggingface.co/nvidia/foundationpose (ONNX, converted to
+PyTorch at startup with onnx2torch).
+Reference implementation: https://github.com/NVlabs/FoundationPose
+"""
+
+import spaces  # noqa: F401  (must be imported before torch / any CUDA touch)
+
+import json
+import math
+import os
+import time
+import traceback
+
+import cv2
+import gradio as gr
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import kornia  # noqa: E402
+import nvdiffrast.torch as dr  # noqa: E402
+from onnx2torch import convert  # noqa: E402
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+EX_DIR = os.path.join(HERE, "examples")
+REPO_ID = "nvidia/foundationpose"
+
+# ---------------------------------------------------------------------------
+# Constants taken verbatim from the two official FoundationPose config.yml files
+# (refiner: 2023-10-28-18-33-37, scorer: 2024-01-11-20-02-45).
+# ---------------------------------------------------------------------------
+INPUT_RESIZE = (160, 160)
+REFINE_CROP_RATIO = 1.2
+SCORE_CROP_RATIO = 1.1
+ROT_NORMALIZER = 0.3490658503988659  # 20 degrees, in radians
+GLCAM_IN_CVCAM = np.array(
+    [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]], dtype=np.float64
+)
+
+MAX_SIDE = 1280  # inputs larger than this get downscaled (K is rescaled too)
+MAX_TEX_SIZE = 2000
+
+# ---------------------------------------------------------------------------
+# Models — converted from ONNX to PyTorch and moved to CUDA at module scope so
+# that ZeroGPU can pack the weights and stream them in on the first request.
+# ---------------------------------------------------------------------------
+print("Downloading FoundationPose ONNX weights ...", flush=True)
+_refiner_path = hf_hub_download(REPO_ID, "refiner_net.onnx")
+_scorer_path = hf_hub_download(REPO_ID, "score_net.onnx")
+
+print("Converting ONNX -> PyTorch ...", flush=True)
+refine_net = convert(_refiner_path).eval().to("cuda")
+score_net = convert(_scorer_path).eval().to("cuda")
+for _p in list(refine_net.parameters()) + list(score_net.parameters()):
+    _p.requires_grad_(False)
+print(
+    f"refine_net params: {sum(p.numel() for p in refine_net.parameters()):,} | "
+    f"score_net params: {sum(p.numel() for p in score_net.parameters()):,}",
+    flush=True,
+)
+
+_GLCTX = None
+
+
+def get_glctx():
+    """nvdiffrast CUDA raster context — created lazily inside the GPU worker."""
+    global _GLCTX
+    if _GLCTX is None:
+        _GLCTX = dr.RasterizeCudaContext()
+    return _GLCTX
+
+
+# ---------------------------------------------------------------------------
+# Geometry / rendering helpers (ports of FoundationPose/Utils.py)
+# ---------------------------------------------------------------------------
+def to_homo_torch(pts):
+    ones = torch.ones((*pts.shape[:-1], 1), dtype=pts.dtype, device=pts.device)
+    return torch.cat((pts, ones), dim=-1)
+
+
+def transform_pts(pts, tf):
+    if len(tf.shape) >= 3 and tf.shape[-3] != pts.shape[-2]:
+        tf = tf[..., None, :, :]
+    return (tf[..., :-1, :-1] @ pts[..., None] + tf[..., :-1, -1:])[..., 0]
+
+
+def transform_dirs(dirs, tf):
+    if len(tf.shape) >= 3 and tf.shape[-3] != dirs.shape[-2]:
+        tf = tf[..., None, :, :]
+    return (tf[..., :3, :3] @ dirs[..., None])[..., 0]
+
+
+def so3_exp_map(log_rot, eps=1e-4):
+    """Rodrigues' formula — matches pytorch3d.transforms.so3_exp_map."""
+    nrms = (log_rot * log_rot).sum(1)
+    rot_angles = torch.clamp(nrms, eps).sqrt()
+    rot_angles_inv = 1.0 / rot_angles
+    fac1 = rot_angles_inv * rot_angles.sin()
+    fac2 = rot_angles_inv * rot_angles_inv * (1.0 - rot_angles.cos())
+    skews = torch.zeros(
+        (log_rot.shape[0], 3, 3), dtype=log_rot.dtype, device=log_rot.device
+    )
+    skews[:, 0, 1] = -log_rot[:, 2]
+    skews[:, 0, 2] = log_rot[:, 1]
+    skews[:, 1, 0] = log_rot[:, 2]
+    skews[:, 1, 2] = -log_rot[:, 0]
+    skews[:, 2, 0] = -log_rot[:, 1]
+    skews[:, 2, 1] = log_rot[:, 0]
+    skews_sq = torch.bmm(skews, skews)
+    eye = torch.eye(3, dtype=log_rot.dtype, device=log_rot.device)[None]
+    return fac1[:, None, None] * skews + fac2[:, None, None] * skews_sq + eye
+
+
+def egocentric_delta_pose_to_pose(a_in_cam, trans_delta, rot_mat_delta):
+    b = torch.eye(4, dtype=torch.float, device=a_in_cam.device)[None].repeat(
+        len(a_in_cam), 1, 1
+    )
+    b[:, :3, 3] = a_in_cam[:, :3, 3] + trans_delta
+    b[:, :3, :3] = rot_mat_delta @ a_in_cam[:, :3, :3]
+    return b
+
+
+def projection_matrix_from_intrinsics(K, height, width, znear, zfar):
+    """Hartley-Zisserman K -> OpenGL projection matrix ('y_down' convention)."""
+    w, h, nc, fc = width, height, znear, zfar
+    depth = float(fc - nc)
+    q = -(fc + nc) / depth
+    qn = -2 * (fc * nc) / depth
+    return np.array(
+        [
+            [2 * K[0, 0] / w, -2 * K[0, 1] / w, (-2 * K[0, 2] + w) / w, 0],
+            [0, 2 * K[1, 1] / h, (2 * K[1, 2] - h) / h, 0],
+            [0, 0, q, qn],
+            [0, 0, -1, 0],
+        ]
+    )
+
+
+def depth2xyzmap_t(depth, K):
+    """depth (H,W) torch -> xyz map (H,W,3) torch, in camera frame."""
+    H, W = depth.shape[-2:]
+    vs, us = torch.meshgrid(
+        torch.arange(H, device=depth.device, dtype=torch.float),
+        torch.arange(W, device=depth.device, dtype=torch.float),
+        indexing="ij",
+    )
+    xs = (us - K[0, 2]) * depth / K[0, 0]
+    ys = (vs - K[1, 2]) * depth / K[1, 1]
+    xyz = torch.stack([xs, ys, depth], dim=-1)
+    xyz[depth < 0.001] = 0
+    return xyz
+
+
+def depth2xyzmap_batch_t(depths, K):
+    """depths (B,H,W) -> (B,H,W,3)."""
+    B, H, W = depths.shape
+    vs, us = torch.meshgrid(
+        torch.arange(H, device=depths.device, dtype=torch.float),
+        torch.arange(W, device=depths.device, dtype=torch.float),
+        indexing="ij",
+    )
+    xs = (us[None] - K[0, 2]) * depths / K[0, 0]
+    ys = (vs[None] - K[1, 2]) * depths / K[1, 1]
+    xyz = torch.stack([xs, ys, depths], dim=-1)
+    xyz[depths < 0.001] = 0
+    return xyz
+
+
+def _unfold(depth, radius):
+    """(H,W) -> patches (K,H,W) and an in-bounds mask (K,H,W)."""
+    k = 2 * radius + 1
+    d = depth[None, None]
+    patches = F.unfold(d, kernel_size=k, padding=radius).reshape(k * k, *depth.shape)
+    inb = F.unfold(
+        torch.ones_like(d), kernel_size=k, padding=radius
+    ).reshape(k * k, *depth.shape) > 0.5
+    return patches, inb
+
+
+def erode_depth(depth, radius=2, depth_diff_thres=0.001, ratio_thres=0.8, zfar=100.0):
+    """Pure-torch port of FoundationPose's warp erode_depth kernel."""
+    patches, inb = _unfold(depth, radius)
+    bad = (patches < 0.001) | (patches >= zfar) | ((patches - depth[None]).abs() > depth_diff_thres)
+    bad_cnt = (bad & inb).sum(0).float()
+    total = inb.sum(0).float().clamp(min=1)
+    out = torch.where(bad_cnt / total > ratio_thres, torch.zeros_like(depth), depth)
+    out = torch.where((depth < 0.001) | (depth >= zfar), torch.zeros_like(depth), out)
+    return out
+
+
+def bilateral_filter_depth(depth, radius=2, zfar=100.0, sigmaD=2.0, sigmaR=100000.0):
+    """Pure-torch port of FoundationPose's warp bilateral_filter_depth kernel."""
+    k = 2 * radius + 1
+    patches, inb = _unfold(depth, radius)
+    valid = (patches >= 0.001) & (patches < zfar) & inb
+    num_valid = valid.sum(0).float()
+    mean_depth = (patches * valid).sum(0) / num_valid.clamp(min=1)
+
+    dv, du = torch.meshgrid(
+        torch.arange(-radius, radius + 1, device=depth.device, dtype=torch.float),
+        torch.arange(-radius, radius + 1, device=depth.device, dtype=torch.float),
+        indexing="ij",
+    )
+    # kernel loop order is u (cols) outer, v (rows) inner -> index = du*k + dv
+    # F.unfold orders the k*k channels row-major: index = (dv+r)*k + (du+r)
+    spatial = torch.exp(-(du * du + dv * dv) / (2.0 * sigmaD * sigmaD)).reshape(k * k, 1, 1)
+
+    sel = valid & ((patches - mean_depth[None]).abs() < 0.01)
+    rng = torch.exp(-((depth[None] - patches) ** 2) / (2.0 * sigmaR * sigmaR))
+    w = spatial * rng * sel
+    sum_w = w.sum(0)
+    out = torch.where(
+        (num_valid > 0) & (sum_w > 0), (w * patches).sum(0) / sum_w.clamp(min=1e-12), torch.zeros_like(depth)
+    )
+    return out
+
+
+def sample_views_icosphere(n_views=40):
+    sub = 1
+    while True:
+        m = trimesh.creation.icosphere(subdivisions=sub, radius=1)
+        if m.vertices.shape[0] >= n_views:
+            break
+        sub += 1
+    cam_in_obs = np.tile(np.eye(4)[None], (len(m.vertices), 1, 1))
+    cam_in_obs[:, :3, 3] = m.vertices
+    up = np.array([0, 0, 1.0])
+    z = -cam_in_obs[:, :3, 3].copy()
+    z /= np.linalg.norm(z, axis=-1, keepdims=True)
+    x = np.cross(up.reshape(1, 3), z)
+    x[(x == 0).all(axis=-1)] = [1, 0, 0]
+    x /= np.linalg.norm(x, axis=-1, keepdims=True)
+    y = np.cross(z, x)
+    y /= np.linalg.norm(y, axis=-1, keepdims=True)
+    cam_in_obs[:, :3, 0] = x
+    cam_in_obs[:, :3, 1] = y
+    cam_in_obs[:, :3, 2] = z
+    return cam_in_obs
+
+
+def make_rotation_grid(min_n_views=40, inplane_step=60):
+    """252 pose hypotheses: 42 icosphere viewpoints x 6 in-plane rotations.
+
+    The reference additionally calls mycpp.cluster_poses(30deg, ...); the minimum
+    pairwise geodesic distance in this grid is 31.7deg so that call is a no-op,
+    which matches the 252-row output documented on the model card.
+    """
+    cam_in_obs = sample_views_icosphere(min_n_views)
+    grid = []
+    for i in range(len(cam_in_obs)):
+        for ang in np.deg2rad(np.arange(0, 360, inplane_step)):
+            rz = np.eye(4)
+            c, s = np.cos(ang), np.sin(ang)
+            rz[0, 0], rz[0, 1], rz[1, 0], rz[1, 1] = c, -s, s, c
+            grid.append(np.linalg.inv(cam_in_obs[i] @ rz))
+    return np.asarray(grid)
+
+
+ROT_GRID = make_rotation_grid()
+print(f"rotation grid: {ROT_GRID.shape}", flush=True)
+
+
+def compute_crop_window_tf_batch(poses, K, crop_ratio, out_size, mesh_diameter):
+    """box_3d crop: a square window around the projected object centre."""
+    B = len(poses)
+    r = mesh_diameter * crop_ratio / 2
+    offsets = torch.tensor(
+        [[0, 0, 0], [r, 0, 0], [-r, 0, 0], [0, r, 0], [0, -r, 0]],
+        device=poses.device,
+        dtype=torch.float,
+    )
+    pts = poses[:, :3, 3].reshape(-1, 1, 3) + offsets.reshape(1, -1, 3)
+    Kt = torch.as_tensor(K, device=poses.device, dtype=torch.float)
+    projected = (Kt @ pts.reshape(-1, 3).T).T
+    uvs = (projected[:, :2] / projected[:, 2:3]).reshape(B, -1, 2)
+    center = uvs[:, 0]
+    rad = torch.abs(uvs - center.reshape(-1, 1, 2)).reshape(B, -1).max(dim=-1)[0]
+    left = (center[:, 0] - rad).round()
+    right = (center[:, 0] + rad).round()
+    top = (center[:, 1] - rad).round()
+    bottom = (center[:, 1] + rad).round()
+    tf = torch.eye(3, device=poses.device, dtype=torch.float)[None].repeat(B, 1, 1)
+    tf[:, 0, 2] = -left
+    tf[:, 1, 2] = -top
+    new_tf = torch.eye(3, device=poses.device, dtype=torch.float)[None].repeat(B, 1, 1)
+    new_tf[:, 0, 0] = out_size[0] / (right - left).clamp(min=1)
+    new_tf[:, 1, 1] = out_size[1] / (bottom - top).clamp(min=1)
+    return new_tf @ tf
+
+
+def _material_image(material):
+    for attr in ("baseColorTexture", "image", "emissiveTexture"):
+        img = getattr(material, attr, None)
+        if img is not None:
+            return img
+    return None
+
+
+def make_mesh_tensors(mesh, max_tex_size=MAX_TEX_SIZE):
+    t = {}
+    img = None
+    if isinstance(mesh.visual, trimesh.visual.texture.TextureVisuals):
+        img = _material_image(mesh.visual.material)
+    if img is not None and getattr(mesh.visual, "uv", None) is not None:
+        arr = np.array(img.convert("RGB"))[..., :3]
+        big = max(arr.shape[0], arr.shape[1])
+        if big > max_tex_size:
+            s = max_tex_size / big
+            arr = cv2.resize(arr, fx=s, fy=s, dsize=None)
+        t["tex"] = torch.as_tensor(
+            np.ascontiguousarray(arr), device="cuda", dtype=torch.float
+        )[None] / 255.0
+        t["uv_idx"] = torch.as_tensor(mesh.faces, device="cuda", dtype=torch.int)
+        uv = torch.as_tensor(np.asarray(mesh.visual.uv), device="cuda", dtype=torch.float).clone()
+        uv[:, 1] = 1 - uv[:, 1]
+        t["uv"] = uv
+    else:
+        vc = None
+        try:
+            vc = np.asarray(mesh.visual.to_color().vertex_colors)
+        except Exception:
+            pass
+        if vc is None or len(vc) != len(mesh.vertices):
+            vc = np.tile(np.array([[160, 160, 160, 255]]), (len(mesh.vertices), 1))
+        t["vertex_color"] = torch.as_tensor(
+            vc[..., :3], device="cuda", dtype=torch.float
+        ) / 255.0
+    t["pos"] = torch.tensor(np.asarray(mesh.vertices), device="cuda", dtype=torch.float)
+    t["faces"] = torch.tensor(np.asarray(mesh.faces), device="cuda", dtype=torch.int)
+    t["vnormals"] = torch.tensor(
+        np.asarray(mesh.vertex_normals), device="cuda", dtype=torch.float
+    )
+    return t
+
+
+def nvdiffrast_render(
+    K, H, W, ob_in_cams, mesh_tensors, output_size=None, bbox2d=None, use_light=True, extra=None
+):
+    glctx = get_glctx()
+    pos = mesh_tensors["pos"]
+    pos_idx = mesh_tensors["faces"]
+    has_tex = "tex" in mesh_tensors
+
+    glcam = torch.tensor(GLCAM_IN_CVCAM, device="cuda", dtype=torch.float)[None]
+    ob_in_glcams = glcam @ ob_in_cams
+    proj = projection_matrix_from_intrinsics(K, height=H, width=W, znear=0.001, zfar=100)
+    proj = torch.as_tensor(proj.reshape(-1, 4, 4), device="cuda", dtype=torch.float)
+    mtx = proj @ ob_in_glcams
+
+    if output_size is None:
+        output_size = np.asarray([H, W])
+
+    pts_cam = transform_pts(pos, ob_in_cams)
+    pos_homo = to_homo_torch(pos)
+    pos_clip = (mtx[:, None] @ pos_homo[None, ..., None])[..., 0]
+    if bbox2d is not None:
+        l = bbox2d[:, 0]
+        t_ = H - bbox2d[:, 1]
+        r = bbox2d[:, 2]
+        b = H - bbox2d[:, 3]
+        tf = torch.eye(4, dtype=torch.float, device="cuda")[None].repeat(len(ob_in_cams), 1, 1)
+        tf[:, 0, 0] = W / (r - l)
+        tf[:, 1, 1] = H / (t_ - b)
+        tf[:, 3, 0] = (W - r - l) / (r - l)
+        tf[:, 3, 1] = (H - t_ - b) / (t_ - b)
+        pos_clip = pos_clip @ tf
+
+    rast_out, _ = dr.rasterize(
+        glctx, pos_clip, pos_idx, resolution=np.asarray(output_size, dtype=np.int64)
+    )
+    xyz_map, _ = dr.interpolate(pts_cam, rast_out, pos_idx)
+    depth = xyz_map[..., 2]
+    if has_tex:
+        texc, _ = dr.interpolate(mesh_tensors["uv"], rast_out, mesh_tensors["uv_idx"])
+        color = dr.texture(mesh_tensors["tex"], texc, filter_mode="linear")
+    else:
+        color, _ = dr.interpolate(mesh_tensors["vertex_color"], rast_out, pos_idx)
+
+    if use_light:
+        vnormals_cam = transform_dirs(mesh_tensors["vnormals"], ob_in_cams)
+        light_dir_neg = -torch.as_tensor(
+            np.array([0, 0, 1.0]), dtype=torch.float, device="cuda"
+        )
+        diffuse = (
+            (F.normalize(vnormals_cam, dim=-1) * F.normalize(light_dir_neg, dim=-1))
+            .sum(dim=-1)
+            .clip(0, 1)[..., None]
+        )
+        diffuse_map, _ = dr.interpolate(diffuse, rast_out, pos_idx)
+        color = color * 0.8 + diffuse_map * color * 0.5
+
+    color = color.clip(0, 1)
+    color = color * torch.clamp(rast_out[..., -1:], 0, 1)
+    color = torch.flip(color, dims=[1])
+    depth = torch.flip(depth, dims=[1])
+    if extra is not None:
+        extra["xyz_map"] = torch.flip(xyz_map, dims=[1])
+    return color, depth
+
+
+# ---------------------------------------------------------------------------
+# Refiner / scorer
+# ---------------------------------------------------------------------------
+def _warp_chunked(src, tfs, dsize, mode):
+    """Warp one (C,H,W) source through B different homographies, in chunks.
+
+    Expanding the source to (B,C,H,W) up front costs ~1 GB at B=252 for a VGA
+    frame; chunking keeps the transient under ~256 MB.
+    """
+    C, H, W = src.shape
+    chunk = max(1, int(2.5e8 / max(C * H * W * 4, 1)))
+    outs = []
+    for b in range(0, len(tfs), chunk):
+        n = len(tfs[b : b + chunk])
+        outs.append(
+            kornia.geometry.transform.warp_perspective(
+                src[None].expand(n, -1, -1, -1).contiguous(),
+                tfs[b : b + chunk],
+                dsize=dsize,
+                mode=mode,
+                align_corners=False,
+            )
+        )
+    return torch.cat(outs, dim=0)
+
+
+def _render_hypotheses(poses, mesh_tensors, K, H, W, tf_to_crops, chunk=128):
+    """Render each hypothesis directly into its own 160x160 crop window."""
+    bbox2d_crop = torch.as_tensor(
+        np.array([0, 0, INPUT_RESIZE[0] - 1, INPUT_RESIZE[1] - 1]).reshape(2, 2),
+        device="cuda",
+        dtype=torch.float,
+    )
+    bbox2d_ori = transform_pts(bbox2d_crop, tf_to_crops.inverse()).reshape(-1, 4)
+    rgb_rs, xyz_rs = [], []
+    for b in range(0, len(poses), chunk):
+        extra = {}
+        rgb_r, _ = nvdiffrast_render(
+            K=K,
+            H=H,
+            W=W,
+            ob_in_cams=poses[b : b + chunk],
+            mesh_tensors=mesh_tensors,
+            output_size=INPUT_RESIZE,
+            bbox2d=bbox2d_ori[b : b + chunk],
+            use_light=True,
+            extra=extra,
+        )
+        rgb_rs.append(rgb_r)
+        xyz_rs.append(extra["xyz_map"])
+    rgb_rs = torch.cat(rgb_rs, dim=0).permute(0, 3, 1, 2) * 255
+    xyz_rs = torch.cat(xyz_rs, dim=0).permute(0, 3, 1, 2)
+    return rgb_rs, xyz_rs
+
+
+def _normalize_xyz(xyz, pose_t, mesh_radius, z_thres):
+    """FoundationPose transform_depth_to_xyzmap (normalize_xyz=True branch)."""
+    bs = xyz.shape[0]
+    invalid = xyz[:, 2:3] < z_thres
+    xyz = xyz - pose_t.reshape(bs, 3, 1, 1)
+    xyz = xyz * (1.0 / mesh_radius)
+    invalid = invalid.expand(bs, 3, -1, -1) | (torch.abs(xyz) >= 2)
+    xyz = xyz.masked_fill(invalid, 0)
+    return xyz
+
+
+@torch.inference_mode()
+def refine_poses(poses, mesh_tensors, rgb_t, xyz_map_t, K, mesh_diameter, iterations, net_bs=512):
+    H, W = rgb_t.shape[:2]
+    B_in_cams = poses
+    for _ in range(iterations):
+        tf_to_crops = compute_crop_window_tf_batch(
+            B_in_cams, K, REFINE_CROP_RATIO, INPUT_RESIZE, mesh_diameter
+        )
+        rgb_as, xyz_as = _render_hypotheses(B_in_cams, mesh_tensors, K, H, W, tf_to_crops)
+        B = len(B_in_cams)
+        rgb_bs = _warp_chunked(
+            rgb_t.permute(2, 0, 1).contiguous(), tf_to_crops, INPUT_RESIZE, "bilinear"
+        )
+        xyz_bs = _warp_chunked(
+            xyz_map_t.permute(2, 0, 1).contiguous(), tf_to_crops, INPUT_RESIZE, "nearest"
+        )
+        rgb_as = rgb_as / 255.0
+        rgb_bs = rgb_bs / 255.0
+        pose_t = B_in_cams[:, :3, 3]
+        radius = mesh_diameter / 2
+        xyz_as = _normalize_xyz(xyz_as, pose_t, radius, 0.001)
+        xyz_bs = _normalize_xyz(xyz_bs, pose_t, radius, 0.001)
+
+        out_poses = []
+        for b in range(0, B, net_bs):
+            A = torch.cat([rgb_as[b : b + net_bs], xyz_as[b : b + net_bs]], dim=1).float()
+            Bt = torch.cat([rgb_bs[b : b + net_bs], xyz_bs[b : b + net_bs]], dim=1).float()
+            trans, rot = refine_net(A, Bt)
+            trans = trans.float()
+            rot = rot.float()
+            trans_delta = trans * (mesh_diameter / 2)
+            rot_mat_delta = torch.tanh(rot) * ROT_NORMALIZER
+            rot_mat_delta = so3_exp_map(rot_mat_delta).permute(0, 2, 1)
+            out_poses.append(
+                egocentric_delta_pose_to_pose(
+                    B_in_cams[b : b + net_bs], trans_delta, rot_mat_delta
+                )
+            )
+        B_in_cams = torch.cat(out_poses, dim=0).reshape(-1, 4, 4)
+    return B_in_cams
+
+
+@torch.inference_mode()
+def score_poses(poses, mesh_tensors, rgb_t, depth_t, K, mesh_diameter, net_bs=512):
+    H, W = rgb_t.shape[:2]
+    B = len(poses)
+    tf_to_crops = compute_crop_window_tf_batch(
+        poses, K, SCORE_CROP_RATIO, INPUT_RESIZE, mesh_diameter
+    )
+    rgb_as, xyz_as = _render_hypotheses(poses, mesh_tensors, K, H, W, tf_to_crops)
+    rgb_bs = _warp_chunked(
+        rgb_t.permute(2, 0, 1).contiguous(), tf_to_crops, INPUT_RESIZE, "bilinear"
+    )
+    # The scorer reconstructs xyzB by warping the cropped depth back to full
+    # resolution, back-projecting there, then warping into the crop again.
+    crop_to_oris = tf_to_crops.inverse()
+    depth_bs = _warp_chunked(
+        depth_t[None].contiguous(), tf_to_crops, INPUT_RESIZE, "nearest"
+    )
+    chunk = max(1, int(1.2e8 / max(H * W, 1)))
+    xyz_bs = []
+    for b in range(0, B, chunk):
+        d_ori = kornia.geometry.transform.warp_perspective(
+            depth_bs[b : b + chunk],
+            crop_to_oris[b : b + chunk],
+            dsize=(H, W),
+            mode="nearest",
+            align_corners=False,
+        )
+        xyz_ori = depth2xyzmap_batch_t(d_ori[:, 0], K).permute(0, 3, 1, 2)
+        xyz_bs.append(
+            kornia.geometry.transform.warp_perspective(
+                xyz_ori,
+                tf_to_crops[b : b + chunk],
+                dsize=INPUT_RESIZE,
+                mode="nearest",
+                align_corners=False,
+            )
+        )
+        del d_ori, xyz_ori
+    xyz_bs = torch.cat(xyz_bs, dim=0)
+
+    rgb_as = rgb_as / 255.0
+    rgb_bs = rgb_bs / 255.0
+    pose_t = poses[:, :3, 3]
+    radius = mesh_diameter / 2
+    xyz_as = _normalize_xyz(xyz_as, pose_t, radius, 0.1)
+    xyz_bs = _normalize_xyz(xyz_bs, pose_t, radius, 0.1)
+
+    scores = []
+    for b in range(0, B, net_bs):
+        A = torch.cat([rgb_as[b : b + net_bs], xyz_as[b : b + net_bs]], dim=1).float()
+        Bt = torch.cat([rgb_bs[b : b + net_bs], xyz_bs[b : b + net_bs]], dim=1).float()
+        scores.append(score_net(A, Bt).float().reshape(-1))
+    return torch.cat(scores, dim=0)
+
+
+# ---------------------------------------------------------------------------
+# Drawing
+# ---------------------------------------------------------------------------
+def _project(pt, K, ob_in_cam):
+    p = K @ ((ob_in_cam @ pt.reshape(4, 1))[:3, :])
+    p = p.reshape(-1) / p.reshape(-1)[2]
+    return tuple(np.round(p[:2]).astype(int).tolist())
+
+
+def draw_xyz_axis(img_rgb, ob_in_cam, K, scale=0.1, thickness=3):
+    img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+    o = _project(np.array([0, 0, 0, 1.0]), K, ob_in_cam)
+    for vec, col in (
+        (np.array([scale, 0, 0, 1.0]), (0, 0, 255)),
+        (np.array([0, scale, 0, 1.0]), (0, 255, 0)),
+        (np.array([0, 0, scale, 1.0]), (255, 0, 0)),
+    ):
+        img = cv2.arrowedLine(
+            img, o, _project(vec, K, ob_in_cam), color=col, thickness=thickness,
+            line_type=cv2.LINE_AA, tipLength=0.15,
+        )
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+
+def draw_posed_3d_box(K, img, ob_in_cam, bbox, line_color=(0, 255, 0), linewidth=2):
+    xmin, ymin, zmin = bbox.min(axis=0)
+    xmax, ymax, zmax = bbox.max(axis=0)
+
+    def line3d(start, end, im):
+        pts = np.stack((start, end), axis=0).reshape(-1, 3)
+        pts = (ob_in_cam @ np.concatenate([pts, np.ones((2, 1))], axis=-1).T).T[:, :3]
+        pr = (K @ pts.T).T
+        uv = np.round(pr[:, :2] / pr[:, 2].reshape(-1, 1)).astype(int)
+        return cv2.line(
+            im, uv[0].tolist(), uv[1].tolist(), color=line_color,
+            thickness=linewidth, lineType=cv2.LINE_AA,
+        )
+
+    for y in [ymin, ymax]:
+        for z in [zmin, zmax]:
+            img = line3d(np.array([xmin, y, z]), np.array([xmax, y, z]), img)
+    for x in [xmin, xmax]:
+        for z in [zmin, zmax]:
+            img = line3d(np.array([x, ymin, z]), np.array([x, ymax, z]), img)
+    for x in [xmin, xmax]:
+        for y in [ymin, ymax]:
+            img = line3d(np.array([x, y, zmin]), np.array([x, y, zmax]), img)
+    return img
+
+
+# ---------------------------------------------------------------------------
+# Input parsing
+# ---------------------------------------------------------------------------
+def load_depth(path, unit):
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".npy":
+        d = np.load(path).astype(np.float32)
+    else:
+        d = np.array(Image.open(path)).astype(np.float32)
+    if d.ndim == 3:
+        d = d[..., 0]
+    scale = {"millimetres (uint16 PNG)": 1e-3, "metres (float)": 1.0, "0.1 mm": 1e-4}[unit]
+    return d * scale
+
+
+def load_mesh(path, unit):
+    obj = trimesh.load(path, process=False, force="mesh")
+    if isinstance(obj, trimesh.Scene):
+        obj = obj.to_geometry()
+    if not isinstance(obj, trimesh.Trimesh):
+        raise gr.Error("Could not read a triangle mesh from the uploaded CAD model.")
+    ext = float(np.linalg.norm(obj.extents))
+    if unit == "auto":
+        if ext > 10.0:
+            obj.vertices = np.asarray(obj.vertices) * 1e-3
+            note = "auto-detected millimetres"
+        else:
+            note = "auto-detected metres"
+    elif unit == "millimetres":
+        obj.vertices = np.asarray(obj.vertices) * 1e-3
+        note = "millimetres"
+    else:
+        note = "metres"
+    return obj, note
+
+
+def parse_floats(text, n, name):
+    try:
+        vals = [float(x) for x in str(text).replace(";", ",").replace(" ", ",").split(",") if x != ""]
+    except ValueError:
+        raise gr.Error(f"Could not parse {name}: {text!r}")
+    if len(vals) != n:
+        raise gr.Error(f"{name} needs {n} numbers, got {len(vals)}: {text!r}")
+    return vals
+
+
+def mesh_diameter(mesh):
+    """Exact diameter via the convex hull (deterministic form of the reference's
+    random-sample pairwise max distance)."""
+    try:
+        pts = np.asarray(mesh.convex_hull.vertices)
+    except Exception:
+        pts = np.asarray(mesh.vertices)
+    if len(pts) > 4000:
+        idx = np.linspace(0, len(pts) - 1, 4000).astype(int)
+        pts = pts[idx]
+    d = np.linalg.norm(pts[None] - pts[:, None], axis=-1)
+    return float(d.max())
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+def _run_pose(
+    rgb_image,
+    depth_file,
+    mesh_file,
+    bbox,
+    intrinsics,
+    depth_unit,
+    mesh_unit,
+    refine_iterations,
+    progress,
+):
+    t_start = time.time()
+    if rgb_image is None:
+        raise gr.Error("Please provide an RGB image.")
+    if depth_file is None:
+        raise gr.Error("Please provide a depth map (16-bit PNG or .npy).")
+    if mesh_file is None:
+        raise gr.Error("Please provide a CAD model of the object.")
+
+    progress(0.05, desc="Reading inputs")
+    rgb = np.array(Image.open(rgb_image).convert("RGB"))
+    depth = load_depth(depth_file, depth_unit)
+    if depth.shape[:2] != rgb.shape[:2]:
+        depth = cv2.resize(depth, (rgb.shape[1], rgb.shape[0]), interpolation=cv2.INTER_NEAREST)
+
+    fx, fy, cx, cy = parse_floats(intrinsics, 4, "intrinsics (fx,fy,cx,cy)")
+    x1, y1, x2, y2 = parse_floats(bbox, 4, "bounding box (x1,y1,x2,y2)")
+    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float64)
+
+    H, W = rgb.shape[:2]
+    if max(H, W) > MAX_SIDE:
+        s = MAX_SIDE / max(H, W)
+        rgb = cv2.resize(rgb, None, fx=s, fy=s, interpolation=cv2.INTER_AREA)
+        depth = cv2.resize(depth, (rgb.shape[1], rgb.shape[0]), interpolation=cv2.INTER_NEAREST)
+        K[:2] *= s
+        x1, y1, x2, y2 = [v * s for v in (x1, y1, x2, y2)]
+        H, W = rgb.shape[:2]
+
+    x1, x2 = sorted([int(round(x1)), int(round(x2))])
+    y1, y2 = sorted([int(round(y1)), int(round(y2))])
+    x1 = max(0, min(W - 2, x1)); x2 = max(x1 + 1, min(W - 1, x2))
+    y1 = max(0, min(H - 2, y1)); y2 = max(y1 + 1, min(H - 1, y2))
+
+    # nvdiffrast's CUDA rasteriser needs both dimensions divisible by 8. Pad on
+    # the right/bottom so pixel coordinates (and therefore K) are unchanged.
+    H0, W0 = H, W
+    Hp, Wp = (H + 7) // 8 * 8, (W + 7) // 8 * 8
+    if (Hp, Wp) != (H, W):
+        rgb = np.pad(rgb, ((0, Hp - H), (0, Wp - W), (0, 0)))
+        depth = np.pad(depth, ((0, Hp - H), (0, Wp - W)))
+        H, W = Hp, Wp
+
+    mesh, unit_note = load_mesh(mesh_file, mesh_unit)
+    progress(0.15, desc="Preparing mesh")
+
+    # Centre the mesh on its AABB centre (reference: FoundationPose.reset_object).
+    model_center = (np.asarray(mesh.vertices).min(axis=0) + np.asarray(mesh.vertices).max(axis=0)) / 2
+    mesh_c = mesh.copy()
+    mesh_c.vertices = np.asarray(mesh_c.vertices) - model_center.reshape(1, 3)
+    diameter = mesh_diameter(mesh_c)
+    if not np.isfinite(diameter) or diameter <= 0:
+        raise gr.Error("Degenerate CAD model (zero diameter).")
+
+    mesh_tensors = make_mesh_tensors(mesh_c)
+
+    rgb_t = torch.as_tensor(rgb.astype(np.float32), device="cuda", dtype=torch.float)
+    depth_t = torch.as_tensor(depth.astype(np.float32), device="cuda", dtype=torch.float)
+
+    progress(0.25, desc="Filtering depth")
+    depth_t = erode_depth(depth_t, radius=2)
+    depth_t = bilateral_filter_depth(depth_t, radius=2)
+
+    # Translation guess from the box centre + median depth inside the box.
+    mask = torch.zeros_like(depth_t, dtype=torch.bool)
+    mask[y1 : y2 + 1, x1 : x2 + 1] = True
+    valid = mask & (depth_t >= 0.001)
+    if int(valid.sum()) < 4:
+        raise gr.Error(
+            "No valid depth inside the box. Check the depth unit and the box coordinates."
+        )
+    zc = torch.median(depth_t[valid]).item()
+    uc = (x1 + x2) / 2.0
+    vc = (y1 + y2) / 2.0
+    center = (np.linalg.inv(K) @ np.array([uc, vc, 1.0]).reshape(3, 1)).reshape(3) * zc
+
+    poses = torch.as_tensor(ROT_GRID, device="cuda", dtype=torch.float).clone()
+    poses[:, :3, 3] = torch.as_tensor(center.reshape(1, 3), device="cuda", dtype=torch.float)
+
+    xyz_map_t = depth2xyzmap_t(depth_t, K)
+
+    progress(0.35, desc=f"Refining {len(poses)} hypotheses x {refine_iterations}")
+    t0 = time.time()
+    poses = refine_poses(
+        poses, mesh_tensors, rgb_t, xyz_map_t, K, diameter, int(refine_iterations)
+    )
+    t_refine = time.time() - t0
+
+    progress(0.8, desc="Scoring hypotheses")
+    t0 = time.time()
+    scores = score_poses(poses, mesh_tensors, rgb_t, depth_t, K, diameter)
+    t_score = time.time() - t0
+
+    order = scores.argsort(descending=True)
+    best = poses[order[0]]
+    best_score = float(scores[order[0]])
+
+    tf_to_centered = np.eye(4)
+    tf_to_centered[:3, 3] = -model_center
+    pose = best.detach().cpu().numpy().astype(np.float64) @ tf_to_centered
+
+    progress(0.9, desc="Rendering")
+    # Overlay: render the mesh at the winning pose over the input image.
+    with torch.inference_mode():
+        color, _ = nvdiffrast_render(
+            K=K,
+            H=H,
+            W=W,
+            ob_in_cams=best[None],
+            mesh_tensors=mesh_tensors,
+            output_size=(H, W),
+            use_light=True,
+        )
+    render = (color[0].clamp(0, 1).cpu().numpy() * 255).astype(np.uint8)
+    alpha = (render.sum(axis=-1) > 0).astype(np.float32)[..., None]
+    overlay = (rgb.astype(np.float32) * (1 - 0.65 * alpha) + render.astype(np.float32) * 0.65 * alpha)
+    overlay = overlay.clip(0, 255).astype(np.uint8)
+
+    # Annotated: oriented 3D bounding box + object axes.
+    to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
+    box = np.stack([-extents / 2, extents / 2], axis=0).reshape(2, 3)
+    center_pose = pose @ np.linalg.inv(to_origin)
+    annotated = rgb.copy()
+    annotated = draw_posed_3d_box(K, annotated, center_pose, box)
+    axis_scale = float(np.clip(0.6 * float(extents.max()), 0.03, 0.15))
+    annotated = draw_xyz_axis(annotated, center_pose, K, scale=axis_scale, thickness=3)
+    annotated = cv2.rectangle(
+        annotated, (x1, y1), (x2, y2), color=(255, 180, 0), thickness=1, lineType=cv2.LINE_AA
+    )
+
+    # Undo the multiple-of-8 padding applied for the rasteriser.
+    annotated = annotated[:H0, :W0]
+    overlay = overlay[:H0, :W0]
+
+    R = pose[:3, :3]
+    t = pose[:3, 3]
+    ang = math.degrees(math.acos(float(np.clip((np.trace(R) - 1) / 2, -1, 1))))
+    pose_txt = "\n".join(
+        "  ".join(f"{v: .6f}" for v in row) for row in pose
+    )
+    total = time.time() - t_start
+    report = (
+        f"**Object → camera translation**  x={t[0]*100:.1f} cm, y={t[1]*100:.1f} cm, z={t[2]*100:.1f} cm\n\n"
+        f"**Rotation angle** {ang:.1f}°  ·  **Score logit** {best_score:.3f}\n\n"
+        f"CAD model: {len(mesh.vertices):,} vertices, diameter {diameter*100:.1f} cm ({unit_note})\n\n"
+        f"252 hypotheses · {int(refine_iterations)} refine passes · "
+        f"refine {t_refine:.2f}s · score {t_score:.2f}s · total {total:.2f}s"
+    )
+
+    del mesh_tensors, rgb_t, depth_t, xyz_map_t, poses, scores
+    torch.cuda.empty_cache()
+    return annotated, overlay, pose_txt, report
+
+
+@spaces.GPU(duration=90)
+def estimate_pose(
+    rgb_image: str,
+    depth_file: str,
+    mesh_file: str,
+    bbox: str,
+    intrinsics: str,
+    depth_unit: str = "millimetres (uint16 PNG)",
+    mesh_unit: str = "auto",
+    refine_iterations: int = 5,
+    progress=gr.Progress(track_tqdm=False),
+):
+    """Estimate the 6-DoF pose of a CAD model in an RGB-D frame with FoundationPose.
+
+    Args:
+        rgb_image: path to the RGB image of the scene.
+        depth_file: path to the aligned depth map (16-bit PNG or .npy).
+        mesh_file: path to the object's CAD model (.glb/.obj/.ply/.stl).
+        bbox: 2D box around the object, "x1,y1,x2,y2" in pixels.
+        intrinsics: pinhole camera intrinsics, "fx,fy,cx,cy" in pixels.
+        depth_unit: unit of the stored depth values.
+        mesh_unit: unit of the CAD model vertices.
+        refine_iterations: number of pose-refinement passes (reference uses 5).
+
+    Returns:
+        Annotated image, rendered overlay, 4x4 object-to-camera pose, and a report.
+    """
+    try:
+        return _run_pose(
+            rgb_image,
+            depth_file,
+            mesh_file,
+            bbox,
+            intrinsics,
+            depth_unit,
+            mesh_unit,
+            refine_iterations,
+            progress,
+        )
+    except gr.Error:
+        raise
+    except Exception as e:
+        traceback.print_exc()
+        raise gr.Error(f"{type(e).__name__}: {e}")
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+_meta_path = os.path.join(EX_DIR, "_meta.json")
+EXAMPLES = []
+if os.path.exists(_meta_path):
+    _meta = json.load(open(_meta_path))
+    for _k in ["mustard_bottle", "power_drill", "pitcher_base", "sugar_box"]:
+        if _k not in _meta:
+            continue
+        m = _meta[_k]
+        EXAMPLES.append(
+            [
+                os.path.join(EX_DIR, m["rgb"]),
+                os.path.join(EX_DIR, m["depth"]),
+                os.path.join(EX_DIR, m["mesh"]),
+                ",".join(str(int(v)) for v in m["bbox"]),
+                ",".join(f"{v:.4f}" for v in m["K"]),
+            ]
+        )
+
+CSS = """
+.dark .gradio-container { --body-background-fill: #06080d; }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="FoundationPose 6-DoF") as demo:
+    gr.Markdown(
+        """
+# FoundationPose — 6-DoF pose of novel objects
+
+Give it an **RGB-D frame**, the object's **CAD model**, the **camera intrinsics**
+and a **2D box**, and [NVIDIA FoundationPose](https://huggingface.co/nvidia/foundationpose)
+returns the object's full 6-DoF pose — zero-shot, no training on the object.
+
+252 pose hypotheses are sampled on an icosphere, refined by the refiner network, then ranked by the scorer network.
+"""
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            rgb_in = gr.Image(type="filepath", label="RGB image", height=280)
+            depth_in = gr.File(
+                label="Depth map — 16-bit PNG or .npy, aligned to the RGB image",
+                file_types=[".png", ".npy", ".tif", ".tiff"],
+            )
+            bbox_in = gr.Textbox(
+                label="2D bounding box — x1,y1,x2,y2 (pixels)",
+                placeholder="409,47,542,324",
+                info="Click twice on the RGB image to set two opposite corners.",
+            )
+            k_in = gr.Textbox(
+                label="Camera intrinsics — fx,fy,cx,cy (pixels)",
+                value="1066.7780,1067.4870,312.9869,241.3109",
+            )
+        with gr.Column(scale=1):
+            mesh_in = gr.Model3D(label="CAD model of the object", height=340)
+            with gr.Accordion("Advanced options", open=False):
+                depth_unit_in = gr.Radio(
+                    ["millimetres (uint16 PNG)", "metres (float)", "0.1 mm"],
+                    value="millimetres (uint16 PNG)",
+                    label="Depth unit",
+                )
+                mesh_unit_in = gr.Radio(
+                    ["auto", "metres", "millimetres"],
+                    value="auto",
+                    label="CAD model unit",
+                )
+                iters_in = gr.Slider(
+                    1, 10, value=5, step=1,
+                    label="Refinement iterations",
+                    info="The reference implementation uses 5.",
+                )
+            run_btn = gr.Button("Estimate pose", variant="primary", size="lg")
+
+    with gr.Row():
+        annotated_out = gr.Image(label="3D box + object axes", height=380)
+        overlay_out = gr.Image(label="CAD model rendered at the estimated pose", height=380)
+    report_out = gr.Markdown()
+    pose_out = gr.Code(label="Object → camera pose (4×4, metres)", language=None)
+
+    corner_state = gr.State(None)
+
+    def on_click(corner, evt: gr.SelectData):
+        x, y = int(evt.index[0]), int(evt.index[1])
+        if corner is None:
+            return (x, y), f"{x},{y},{x},{y}"
+        x0, y0 = corner
+        return None, f"{min(x0,x)},{min(y0,y)},{max(x0,x)},{max(y0,y)}"
+
+    rgb_in.select(on_click, inputs=[corner_state], outputs=[corner_state, bbox_in])
+
+    inputs = [rgb_in, depth_in, mesh_in, bbox_in, k_in, depth_unit_in, mesh_unit_in, iters_in]
+    outputs = [annotated_out, overlay_out, pose_out, report_out]
+    run_btn.click(estimate_pose, inputs=inputs, outputs=outputs)
+
+    if EXAMPLES:
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[rgb_in, depth_in, mesh_in, bbox_in, k_in],
+            outputs=outputs,
+            fn=estimate_pose,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples — YCB-Video (BOP), one of FoundationPose's own evaluation datasets",
+        )
+
+    gr.Markdown(
+        """
+---
+**Model** [nvidia/foundationpose](https://huggingface.co/nvidia/foundationpose) ·
+**Paper** [FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects](https://arxiv.org/abs/2312.08344) (CVPR 2024, Best Paper Nominee) ·
+**Code** [NVlabs/FoundationPose](https://github.com/NVlabs/FoundationPose)
+
+The published checkpoints are ONNX; they are converted to PyTorch with
+[onnx2torch](https://github.com/ENOT-AutoDL/onnx2torch) at startup. The exported
+scorer rates each hypothesis independently (`score_logit` per pose), which is the
+shape NVIDIA ships.
+
+Example scenes and CAD models come from the
+[YCB-Video / BOP](https://huggingface.co/datasets/bop-benchmark/ycbv) dataset (MIT licence,
+© 2017 UW Robotics and State Estimation Lab). Model weights are covered by the
+NVIDIA Open Model License.
+"""
+    )
+
+if __name__ == "__main__":
+    demo.queue(max_size=12).launch(mcp_server=True, show_error=True)

examples/004_sugar_box.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8b3830dd3067ecee4abf16da3624a78353a472b06bc9f2bbaa8f00c500ff65d
+size 1571864

examples/006_mustard_bottle.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9dc011cddcb899fdc8713b4904ca1eb60bda0d0f8e280e726c92de4a247c7a2
+size 1289596

examples/019_pitcher_base.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bda6f4993f9c30be5f052c51ca7671a32d5e96e3baf53359f660bbad44f1f497
+size 1212516

examples/035_power_drill.glb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:777e3f8921db35fbc210c936cc3321b05200c1b318783f804daa414c06347767
+size 1089540

examples/_meta.json

@@ -0,0 +1,186 @@
+{
+ "mustard_bottle": {
+  "name": "006_mustard_bottle",
+  "mesh": "006_mustard_bottle.glb",
+  "rgb": "mustard_bottle_rgb.png",
+  "depth": "mustard_bottle_depth.png",
+  "bbox": [
+   409,
+   47,
+   542,
+   324
+  ],
+  "K": [
+   1066.778,
+   1067.487,
+   312.9869,
+   241.3109
+  ],
+  "gt_pose": [
+   [
+    -0.9586628303517895,
+    -0.2835287329246605,
+    -0.024014888195919924,
+    0.11510576923433374
+   ],
+   [
+    -0.08379337156245308,
+    0.3619577494350433,
+    -0.9284211691506922,
+    -0.041077239594684054
+   ],
+   [
+    0.27192578395746053,
+    -0.8880300419285506,
+    -0.37075363826953395,
+    0.7556857413095329
+   ],
+   [
+    0.0,
+    0.0,
+    0.0,
+    1.0
+   ]
+  ],
+  "scene": "000050",
+  "im": "1874"
+ },
+ "power_drill": {
+  "name": "035_power_drill",
+  "mesh": "035_power_drill.glb",
+  "rgb": "power_drill_rgb.png",
+  "depth": "power_drill_depth.png",
+  "bbox": [
+   58,
+   246,
+   371,
+   456
+  ],
+  "K": [
+   1066.778,
+   1067.487,
+   312.9869,
+   241.3109
+  ],
+  "gt_pose": [
+   [
+    -0.9265112328620186,
+    0.3554015793763549,
+    0.12355990855995728,
+    -0.07007990833197376
+   ],
+   [
+    -0.053990458994419316,
+    -0.4505591780531961,
+    0.8911131272403047,
+    0.06763430848607617
+   ],
+   [
+    0.3723725271017997,
+    0.8189545363840561,
+    0.4366367182365112,
+    0.7316101236716381
+   ],
+   [
+    0.0,
+    0.0,
+    0.0,
+    1.0
+   ]
+  ],
+  "scene": "000050",
+  "im": "1324"
+ },
+ "pitcher_base": {
+  "name": "019_pitcher_base",
+  "mesh": "019_pitcher_base.glb",
+  "rgb": "pitcher_base_rgb.png",
+  "depth": "pitcher_base_depth.png",
+  "bbox": [
+   226,
+   69,
+   513,
+   439
+  ],
+  "K": [
+   1066.778,
+   1067.487,
+   312.9869,
+   241.3109
+  ],
+  "gt_pose": [
+   [
+    -0.7630850378965286,
+    -0.6459310027194527,
+    0.021784917094012195,
+    0.03607193041593472
+   ],
+   [
+    -0.23791608153732055,
+    0.24940775817574404,
+    -0.9387177934933404,
+    0.00816719940180273
+   ],
+   [
+    0.6009142146924277,
+    -0.7215054182161463,
+    -0.3439961790363487,
+    0.7327526085460678
+   ],
+   [
+    0.0,
+    0.0,
+    0.0,
+    1.0
+   ]
+  ],
+  "scene": "000056",
+  "im": "1"
+ },
+ "sugar_box": {
+  "name": "004_sugar_box",
+  "mesh": "004_sugar_box.glb",
+  "rgb": "sugar_box_rgb.png",
+  "depth": "sugar_box_depth.png",
+  "bbox": [
+   151,
+   255,
+   413,
+   452
+  ],
+  "K": [
+   1066.778,
+   1067.487,
+   312.9869,
+   241.3109
+  ],
+  "gt_pose": [
+   [
+    0.04261657197093155,
+    0.8018176555886548,
+    0.5960480236123241,
+    -0.015860798256254367
+   ],
+   [
+    -0.85005714750022,
+    0.3425727880987984,
+    -0.40005896186039525,
+    0.0695873326444807
+   ],
+   [
+    -0.5249637632720354,
+    -0.48962547962401337,
+    0.6961900608197962,
+    0.7274594452418055
+   ],
+   [
+    0.0,
+    0.0,
+    0.0,
+    1.0
+   ]
+  ],
+  "scene": "000051",
+  "im": "1729"
+ }
+}

requirements.txt

@@ -0,0 +1,13 @@
+# torch pinned to match the prebuilt sm_120 nvdiffrast wheel cell (pt211-cu130-cp312)
+torch==2.11.0
+torchvision==0.26.0
+numpy
+scipy
+networkx
+pillow
+opencv-python-headless
+trimesh
+kornia
+onnx
+onnx2torch
+https://huggingface.co/datasets/multimodalart/zerogpu-blackwell-wheels/resolve/main/wheels/pt211-cu130-cp312/nvdiffrast-0.4.0-cp312-cp312-linux_x86_64.whl