app.py

"""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)