Spaces:
Running on Zero
Running on Zero
| """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" | |
| 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 | |
| def predict(image: Image.Image, size: 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) | |
| # Foreground masking is mandatory — keeps the cloud clean and the camera | |
| # marker meaningful (background depth is unreliable). | |
| mask = _foreground_mask(image_pil, h0, w0) | |
| ply_path = _make_ply(image_rgb, pointmap, mask) | |
| return ply_path, ply_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 = """ | |
| <div id="title">Sapiens2: Pointmap</div> | |
| <div id="subtitle">ICLR 2026</div> | |
| <div id="badges"> | |
| <a class="pill" href="https://github.com/facebookresearch/sapiens2" target="_blank" rel="noopener"> | |
| <svg viewBox="0 0 24 24" fill="currentColor"><path d="M12 .3a12 12 0 0 0-3.8 23.4c.6.1.8-.3.8-.6v-2c-3.3.7-4-1.6-4-1.6-.6-1.4-1.4-1.8-1.4-1.8-1.1-.7.1-.7.1-.7 1.3.1 2 1.3 2 1.3 1.1 1.9 3 1.4 3.7 1 .1-.8.4-1.4.8-1.7-2.7-.3-5.5-1.3-5.5-5.9 0-1.3.5-2.4 1.3-3.2-.1-.4-.6-1.6.1-3.2 0 0 1-.3 3.3 1.2a11.5 11.5 0 0 1 6 0c2.3-1.5 3.3-1.2 3.3-1.2.7 1.6.2 2.8.1 3.2.8.8 1.3 1.9 1.3 3.2 0 4.6-2.8 5.6-5.5 5.9.4.4.8 1.1.8 2.2v3.3c0 .3.2.7.8.6A12 12 0 0 0 12 .3"/></svg> | |
| Code | |
| </a> | |
| <a class="pill" href="https://huggingface.co/facebook/sapiens2" target="_blank" rel="noopener"> | |
| 🤗 Models | |
| </a> | |
| <a class="pill" href="https://openreview.net/pdf?id=IVAlYCqdvW" target="_blank" rel="noopener"> | |
| <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M14 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V8z"/><polyline points="14 2 14 8 20 8"/><line x1="9" y1="13" x2="15" y2="13"/><line x1="9" y1="17" x2="15" y2="17"/></svg> | |
| Paper | |
| </a> | |
| <a class="pill" href="https://rawalkhirodkar.github.io/sapiens2" target="_blank" rel="noopener"> | |
| <svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="10"/><line x1="2" y1="12" x2="22" y2="12"/><path d="M12 2a15.3 15.3 0 0 1 4 10 15.3 15.3 0 0 1-4 10 15.3 15.3 0 0 1-4-10 15.3 15.3 0 0 1 4-10z"/></svg> | |
| Project | |
| </a> | |
| </div> | |
| """ | |
| 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=4, | |
| ) | |
| 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)") | |
| run.click(predict, inputs=[inp, size], outputs=[out_ply, out_ply_file]) | |
| 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) | |