"""Sapiens2 pointmap Gradio Space. Image → per-pixel 3D pointmap (camera frame, metric units). The result is exported as a .ply point cloud and rendered with Gradio's Model3D component for interactive 3D viewing. """ import sys import os sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import tempfile import cv2 import gradio as gr import numpy as np import open3d as o3d import spaces import torch import torch.nn.functional as F from PIL import Image from huggingface_hub import hf_hub_download from sapiens.dense.models import PointmapEstimator, init_model # registers in registry _ = PointmapEstimator # ----------------------------------------------------------------------------- # Config ASSETS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets") CONFIGS_DIR = os.path.join(ASSETS_DIR, "configs") POINTMAP_MODELS = { "0.4B": { "repo": "facebook/sapiens2-pointmap-0.4b", "filename": "sapiens2_0.4b_pointmap.safetensors", "config": os.path.join(CONFIGS_DIR, "sapiens2_0.4b_pointmap_render_people-1024x768.py"), }, "0.8B": { "repo": "facebook/sapiens2-pointmap-0.8b", "filename": "sapiens2_0.8b_pointmap.safetensors", "config": os.path.join(CONFIGS_DIR, "sapiens2_0.8b_pointmap_render_people-1024x768.py"), }, "1B": { "repo": "facebook/sapiens2-pointmap-1b", "filename": "sapiens2_1b_pointmap.safetensors", "config": os.path.join(CONFIGS_DIR, "sapiens2_1b_pointmap_render_people-1024x768.py"), }, "5B": { "repo": "facebook/sapiens2-pointmap-5b", "filename": "sapiens2_5b_pointmap.safetensors", "config": os.path.join(CONFIGS_DIR, "sapiens2_5b_pointmap_render_people-1024x768.py"), }, } DEFAULT_SIZE = "1B" DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # ----------------------------------------------------------------------------- # Model cache _pointmap_model_cache: dict = {} def _get_pointmap_model(size: str): if size not in _pointmap_model_cache: spec = POINTMAP_MODELS[size] ckpt = hf_hub_download(repo_id=spec["repo"], filename=spec["filename"]) model = init_model(spec["config"], ckpt, device=DEVICE) _pointmap_model_cache[size] = model return _pointmap_model_cache[size] print("[startup] pre-loading all pointmap sizes ...") for _size in POINTMAP_MODELS: _get_pointmap_model(_size) print("[startup] ready.") # ----------------------------------------------------------------------------- # Inference def _estimate_pointmap(image_bgr: np.ndarray, model) -> np.ndarray: h0, w0 = image_bgr.shape[:2] data = model.pipeline(dict(img=image_bgr)) data = model.data_preprocessor(data) inputs, data_samples = data["inputs"], data["data_samples"] if inputs.ndim == 3: inputs = inputs.unsqueeze(0) with torch.no_grad(): pointmap, scale = model(inputs) # (1, 3, H, W), (1, 1) pointmap = pointmap / scale # convert to metric pad = data_samples["meta"]["padding_size"] pad_left, pad_right, pad_top, pad_bottom = pad pointmap = pointmap[ :, :, pad_top : inputs.shape[2] - pad_bottom, pad_left : inputs.shape[3] - pad_right, ] pointmap = F.interpolate(pointmap, size=(h0, w0), mode="bilinear", align_corners=False) return pointmap.squeeze(0).cpu().float().numpy().transpose(1, 2, 0) # (H, W, 3) # ----------------------------------------------------------------------------- # Point cloud export def _make_ply(image_rgb: np.ndarray, pointmap_hwc: np.ndarray, max_points: int = 200_000) -> str: """Subsample, filter to a reasonable depth range, and write a .ply file.""" pts = pointmap_hwc.reshape(-1, 3) cols = (image_rgb.reshape(-1, 3).astype(np.float32) / 255.0) # Drop points with non-finite or extreme depth z = pts[:, 2] finite = np.isfinite(pts).all(axis=1) & (z > 0.05) & (z < 25.0) pts, cols = pts[finite], cols[finite] if len(pts) > max_points: idx = np.random.default_rng(0).choice(len(pts), size=max_points, replace=False) pts, cols = pts[idx], cols[idx] pc = o3d.geometry.PointCloud() pc.points = o3d.utility.Vector3dVector(pts.astype(np.float64)) pc.colors = o3d.utility.Vector3dVector(cols.astype(np.float64)) out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".ply").name o3d.io.write_point_cloud(out_path, pc, write_ascii=False) return out_path # ----------------------------------------------------------------------------- # Gradio handler @spaces.GPU(duration=180) def predict(image: Image.Image, size: str): if image is None: return None, None image_rgb = np.array(image.convert("RGB")) image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR) model = _get_pointmap_model(size) pointmap = _estimate_pointmap(image_bgr, model) # (H, W, 3) metric, camera frame ply_path = _make_ply(image_rgb, pointmap) npy_path = tempfile.NamedTemporaryFile(delete=False, suffix=".npy").name np.save(npy_path, pointmap.astype(np.float32)) return ply_path, npy_path # ----------------------------------------------------------------------------- # UI EXAMPLES = sorted( os.path.join(ASSETS_DIR, "images", n) for n in os.listdir(os.path.join(ASSETS_DIR, "images")) if n.lower().endswith((".jpg", ".jpeg", ".png")) ) with gr.Blocks(title="Sapiens2 Pointmap", theme=gr.themes.Default()) as demo: gr.Markdown( "# Sapiens2: Pointmap Estimation\n" "### ICLR 2026\n" "Per-pixel 3D pointmap (camera frame, metric units). Drag to rotate the result.\n\n" "[Code](https://github.com/facebookresearch/sapiens2) · " "[Models](https://huggingface.co/facebook/sapiens2) · " "[Paper](https://openreview.net/pdf?id=IVAlYCqdvW)" ) with gr.Row(): with gr.Column(): inp = gr.Image(label="Input", type="pil") size = gr.Radio( choices=list(POINTMAP_MODELS.keys()), value=DEFAULT_SIZE, label="Model size", ) run = gr.Button("Run", variant="primary") gr.Examples(examples=EXAMPLES, inputs=inp, examples_per_page=14) with gr.Column(): out_ply = gr.Model3D(label="Point cloud (drag to rotate)", clear_color=[0.05, 0.05, 0.05, 1.0]) out_npy = gr.File(label="Raw pointmap (.npy float32 [H, W, 3] in meters)") run.click(predict, inputs=[inp, size], outputs=[out_ply, out_npy]) if __name__ == "__main__": if torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True demo.launch(share=False)