"""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. Optionally applies a v1 binary fg/bg mask so only
foreground points end up in the cloud.
"""
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 torchvision import transforms
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 = "0.4B" # iteration mode — only this is preloaded; others lazy-load on click
FG_REPO = "facebook/sapiens-seg-foreground-1b-torchscript"
FG_FILENAME = "sapiens_1b_seg_foreground_epoch_8_torchscript.pt2"
BG_OPTIONS = ["fg-bg", "no-bg-removal"]
DEFAULT_BG = "fg-bg"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
_fg_transform = transforms.Compose([
transforms.Resize((1024, 768)),
transforms.ToTensor(),
transforms.Normalize(mean=[123.5 / 255, 116.5 / 255, 103.5 / 255],
std=[58.5 / 255, 57.0 / 255, 57.5 / 255]),
])
# -----------------------------------------------------------------------------
# Model cache
_pointmap_model_cache: dict = {}
_fg_model = None
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]
def _get_fg_model():
global _fg_model
if _fg_model is None:
ckpt = hf_hub_download(repo_id=FG_REPO, filename=FG_FILENAME)
_fg_model = torch.jit.load(ckpt).eval().to(DEVICE)
return _fg_model
# Iteration mode: only preload the default (0.4B) for fast Space boot.
# Re-enable full preload by uncommenting the loop below.
print("[startup] pre-loading 0.4B (iteration mode) + fg/bg ...")
_get_pointmap_model(DEFAULT_SIZE)
_get_fg_model()
# 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)
pointmap = pointmap / scale # → metric units
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)
def _foreground_mask(image_pil: Image.Image, target_h: int, target_w: int) -> np.ndarray:
fg = _get_fg_model()
inputs = _fg_transform(image_pil).unsqueeze(0).to(DEVICE)
with torch.no_grad():
out = fg(inputs)
out = F.interpolate(out, size=(target_h, target_w), mode="bilinear", align_corners=False)
return (out.argmax(dim=1)[0] > 0).cpu().numpy()
# -----------------------------------------------------------------------------
# Point cloud export
def _camera_marker(radius: float = 0.04, n_points: int = 800,
color=(0.20, 0.55, 0.96)) -> o3d.geometry.PointCloud:
"""Small uniformly-blue sphere at the world origin marking the camera.
Manual Fibonacci-sphere sampling — instant, vs Open3D's poisson-disk which
can take seconds per call.
"""
rng = np.random.default_rng(0)
i = np.arange(n_points)
phi = np.arccos(1 - 2 * (i + 0.5) / n_points) # latitude
theta = np.pi * (1 + 5 ** 0.5) * (i + 0.5) # golden-angle longitude
pts = np.stack([
radius * np.sin(phi) * np.cos(theta),
radius * np.sin(phi) * np.sin(theta),
radius * np.cos(phi),
], axis=1)
pc = o3d.geometry.PointCloud()
pc.points = o3d.utility.Vector3dVector(pts.astype(np.float64))
pc.colors = o3d.utility.Vector3dVector(np.tile(color, (n_points, 1)).astype(np.float64))
return pc
def _make_ply(image_rgb: np.ndarray, pointmap_hwc: np.ndarray, mask_hw: np.ndarray | None = None,
max_points: int = 200_000) -> str:
pts = pointmap_hwc.reshape(-1, 3)
cols = (image_rgb.reshape(-1, 3).astype(np.float32) / 255.0)
z = pts[:, 2]
finite = np.isfinite(pts).all(axis=1) & (z > 0.05) & (z < 25.0)
if mask_hw is not None:
finite &= mask_hw.reshape(-1)
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))
# Add the camera marker (blue ball at origin) so users see where the
# observer is in the reconstructed 3D scene.
pc += _camera_marker()
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, bg_mode: str):
if image is None:
return None, None
image_pil = image.convert("RGB")
image_rgb = np.array(image_pil)
image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
h0, w0 = image_rgb.shape[:2]
model = _get_pointmap_model(size)
pointmap = _estimate_pointmap(image_bgr, model)
mask = _foreground_mask(image_pil, h0, w0) if bg_mode == "fg-bg" else None
ply_path = _make_ply(image_rgb, pointmap, mask)
npy_path = tempfile.NamedTemporaryFile(delete=False, suffix=".npy").name
np.save(npy_path, pointmap.astype(np.float32))
return ply_path, 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"))
)
CUSTOM_CSS = """
:root, body, .gradio-container, button, input, select, textarea,
.gradio-container *:not(code):not(pre) {
font-family: "Helvetica Neue", Helvetica, Arial, sans-serif !important;
-webkit-font-smoothing: antialiased;
-moz-osx-font-smoothing: grayscale;
}
#title { text-align: center; font-size: 44px; font-weight: 700;
letter-spacing: -0.01em; margin: 28px 0 4px;
background: linear-gradient(90deg, #1d4ed8 0%, #6d28d9 50%, #be185d 100%);
-webkit-background-clip: text; -webkit-text-fill-color: transparent;
background-clip: text; }
#subtitle { text-align: center; font-size: 12px; color: #64748b;
letter-spacing: 0.18em; margin: 0 0 14px; text-transform: uppercase;
font-weight: 500; }
#badges { display: flex; justify-content: center; flex-wrap: wrap;
gap: 8px; margin: 0 0 32px; }
.pill { display: inline-flex; align-items: center; gap: 6px;
padding: 7px 14px; border-radius: 999px;
background: #f1f5f9; color: #0f172a !important;
font-size: 13px; font-weight: 500; letter-spacing: 0.01em;
text-decoration: none !important; border: 1px solid #e2e8f0;
transition: background 150ms ease, transform 150ms ease, border-color 150ms ease; }
.pill:hover { background: #0f172a; color: #f8fafc !important;
border-color: #0f172a; transform: translateY(-1px); }
.pill svg { width: 14px; height: 14px; }
"""
HEADER_HTML = """
Sapiens2: Pointmap
ICLR 2026
"""
with gr.Blocks(title="Sapiens2 Pointmap", theme=gr.themes.Soft(), css=CUSTOM_CSS) as demo:
gr.HTML(HEADER_HTML)
with gr.Row(equal_height=True):
inp = gr.Image(label="Input", type="pil", height=640)
out_ply = gr.Model3D(
label="Point cloud — drag to rotate, scroll to zoom, shift+drag to pan",
height=640,
clear_color=[0.07, 0.09, 0.13, 1.0], # subtle slate-900 backdrop
display_mode="point_cloud",
zoom_speed=0.7,
pan_speed=0.5,
)
with gr.Row():
size = gr.Radio(
choices=["0.4B"], # iteration mode — re-add other sizes when shipping
value=DEFAULT_SIZE,
label="Model",
scale=2,
)
bg = gr.Radio(
choices=BG_OPTIONS,
value=DEFAULT_BG,
label="Background",
scale=2,
)
run = gr.Button("Run", variant="primary", size="lg", scale=1)
gr.Examples(examples=EXAMPLES, inputs=inp, examples_per_page=14)
with gr.Accordion("Raw Pointmap", open=False):
out_ply_file = gr.File(label="Point cloud (.ply — open in MeshLab/CloudCompare/Blender)")
out_npy = gr.File(label="Raw pointmap (.npy float32 [H, W, 3] in meters)")
run.click(predict, inputs=[inp, size, bg], outputs=[out_ply, out_ply_file, 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)