"""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) # Outline the rendered silhouette so the alignment is easy to judge. cnts, _ = cv2.findContours( (alpha[..., 0] > 0).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE ) overlay = cv2.drawContours(np.ascontiguousarray(overlay), cnts, -1, (0, 255, 120), 2, cv2.LINE_AA) # 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=40) 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(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], api_visibility="private" ) 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( theme=gr.themes.Citrus(), css=CSS, mcp_server=True, show_error=True )