Spaces:
Running on Zero
Running on Zero
| import spaces # must be imported before torch | |
| import base64 | |
| import csv | |
| import sys | |
| import glob | |
| import html | |
| import json | |
| import os | |
| import tempfile | |
| import time | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from huggingface_hub import hf_hub_download | |
| from solv3r import Solv3R, cameras_from_pose_enc, load_images | |
| from solv3r.export import confident_points, sky_mask, write_ply | |
| from solv3r.images import VIDEO_EXTENSIONS, list_images | |
| from solv3r.model import autocast | |
| from solv3r.network import skyseg | |
| MAX_VIEWS = 64 | |
| # The viewer HTML embeds the point cloud; let Gradio's example cache (a CSV file) hold it. | |
| csv.field_size_limit(sys.maxsize) | |
| # ---------------------------------------------------------------- model (module scope, on GPU) | |
| model = Solv3R.from_pretrained("jgkwak/solv3r", device="cuda") | |
| # The sky segmentation ONNX model (used for the optional sky mask) is fetched once at startup. | |
| skyseg.SKYSEG_ONNX_PATH = hf_hub_download("JianyuanWang/skyseg", "skyseg.onnx") | |
| # ---------------------------------------------------------------- viewer | |
| VIEWER = """<!doctype html> | |
| <html><head><meta charset="utf-8"><title>Solv3R</title> | |
| <style>html,body{margin:0;height:100%;overflow:hidden;background:#0b0d10;color:#d9dde2;font:13px system-ui,sans-serif} | |
| #info{position:absolute;left:12px;top:10px;pointer-events:none;opacity:.85} | |
| #ui{position:absolute;right:12px;top:10px;display:flex;gap:14px;align-items:center;background:rgba(0,0,0,.35);padding:4px 10px;border-radius:8px}</style> | |
| </head><body><div id="info"></div> | |
| <div id="ui"><label>point size <input id="size" type="range" min="1" max="40" value="8"></label> | |
| <label><input id="cams" type="checkbox" checked> cameras</label></div> | |
| <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script> | |
| <script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/controls/OrbitControls.js"></script> | |
| <script> | |
| const D = __DATA__; | |
| const raw = atob(D.points), bytes = new Uint8Array(raw.length); | |
| for (let i = 0; i < raw.length; i++) bytes[i] = raw.charCodeAt(i); | |
| const q = new Uint16Array(bytes.buffer, 0, D.n * 3), pos = new Float32Array(D.n * 3); | |
| for (let i = 0; i < D.n * 3; i++) { const k = i % 3; pos[i] = D.lo[k] + (q[i] / 65535) * (D.hi[k] - D.lo[k]); } | |
| const geo = new THREE.BufferGeometry(); | |
| geo.setAttribute("position", new THREE.BufferAttribute(pos, 3)); | |
| geo.setAttribute("color", new THREE.BufferAttribute(new Uint8Array(bytes.buffer, D.n * 6, D.n * 3), 3, true)); | |
| const mat = new THREE.PointsMaterial({ size: 0.008, vertexColors: true, sizeAttenuation: true }); | |
| const scene = new THREE.Scene(), renderer = new THREE.WebGLRenderer({ antialias: true }); | |
| renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2)); document.body.appendChild(renderer.domElement); | |
| scene.add(new THREE.Points(geo, mat)); | |
| const camGroup = new THREE.Group(); | |
| const fr = new THREE.BufferGeometry().setAttribute("position", new THREE.BufferAttribute(new Float32Array(D.frusta), 3)); | |
| fr.setAttribute("color", new THREE.BufferAttribute(new Float32Array(D.frustaColors), 3)); | |
| camGroup.add(new THREE.LineSegments(fr, new THREE.LineBasicMaterial({ vertexColors: true }))); | |
| if (D.path.length > 1) camGroup.add(new THREE.Line( | |
| new THREE.BufferGeometry().setAttribute("position", new THREE.BufferAttribute(new Float32Array(D.path), 3)), | |
| new THREE.LineBasicMaterial({ color: 0x6ea7e3, transparent: true, opacity: 0.5 }))); | |
| scene.add(camGroup); | |
| const camera = new THREE.PerspectiveCamera(50, 1, 0.01, 100), controls = new THREE.OrbitControls(camera, renderer.domElement); | |
| camera.position.set(0, 0.9, 2.4); controls.update(); | |
| const render = () => renderer.render(scene, camera); | |
| const resize = () => { renderer.setSize(innerWidth, innerHeight); camera.aspect = innerWidth / innerHeight; camera.updateProjectionMatrix(); render(); }; | |
| controls.addEventListener("change", render); addEventListener("resize", resize); | |
| document.getElementById("size").oninput = (e) => { mat.size = e.target.value / 1000; render(); }; | |
| document.getElementById("cams").onchange = (e) => { camGroup.visible = e.target.checked; render(); }; | |
| document.getElementById("info").textContent = `${D.n.toLocaleString()} points, ${D.views} cameras. Drag to rotate, scroll to zoom, right-drag to pan.`; | |
| resize(); | |
| </script></body></html> | |
| """ | |
| def _hsv_colors(n): | |
| out = [] | |
| for i in range(n): | |
| h = (i / max(n, 1)) * 0.8 | |
| c = cv2.cvtColor(np.uint8([[[int(h * 180), 200, 255]]]), cv2.COLOR_HSV2RGB)[0, 0] / 255.0 | |
| out.append(c.tolist()) | |
| return out | |
| def build_viewer(xyz, rgb, c2w, K, hw, max_points=200_000): | |
| """Self-contained three.js viewer (points + camera frusta) embedded through an iframe srcdoc.""" | |
| xyz, rgb = np.asarray(xyz, np.float64), np.asarray(rgb, np.uint8) | |
| if len(xyz) > max_points: | |
| idx = np.sort(np.random.default_rng(0).choice(len(xyz), max_points, replace=False)) | |
| xyz, rgb = xyz[idx], rgb[idx] | |
| c = np.median(xyz, 0) | |
| r = max(float(np.percentile(np.linalg.norm(xyz - c, axis=1), 80)), 1e-9) | |
| up, fwd = -c2w[:, :3, 1].mean(0), c2w[:, :3, 2].mean(0) | |
| up = up / np.linalg.norm(up) | |
| fwd = fwd - (fwd @ up) * up | |
| fwd = fwd / np.linalg.norm(fwd) if np.linalg.norm(fwd) > 1e-6 else np.cross(up, [1.0, 0, 0]) | |
| R = np.stack([np.cross(fwd, up), up, -fwd]) # viewer frame: y up, capture looks along -z | |
| to_view = lambda p: (p - c) @ R.T / r # noqa: E731 | |
| V = to_view(xyz) | |
| keep = np.linalg.norm(V, axis=1) < 6 | |
| V, rgb = V[keep], rgb[keep] | |
| lo, hi = V.min(0), V.max(0) | |
| q = np.round((V - lo) / np.maximum(hi - lo, 1e-9) * 65535).astype("<u2") | |
| # camera frusta: 8 segments per camera | |
| H, W = hw | |
| d = 0.12 * r | |
| segs, cols = [], [] | |
| colors = _hsv_colors(len(c2w)) | |
| for i in range(len(c2w)): | |
| fx, fy, cx, cy = K[i, 0, 0], K[i, 1, 1], K[i, 0, 2], K[i, 1, 2] | |
| corners = np.array([[(u - cx) / fx * d, (v - cy) / fy * d, d] for u, v in [(0, 0), (W, 0), (W, H), (0, H)]]) | |
| wc = corners @ c2w[i, :3, :3].T + c2w[i, :3, 3] | |
| o = c2w[i, :3, 3] | |
| pts = [o, wc[0], o, wc[1], o, wc[2], o, wc[3], wc[0], wc[1], wc[1], wc[2], wc[2], wc[3], wc[3], wc[0]] | |
| segs.append(to_view(np.array(pts))) | |
| cols += [colors[i]] * len(pts) | |
| frusta = np.concatenate(segs).round(4).ravel().tolist() if segs else [] | |
| path = to_view(c2w[:, :3, 3]).round(4).ravel().tolist() | |
| data = dict(n=len(V), views=len(c2w), lo=lo.tolist(), hi=hi.tolist(), frusta=frusta, | |
| frustaColors=np.array(cols).round(3).ravel().tolist(), path=path, | |
| points=base64.b64encode(q.tobytes() + np.ascontiguousarray(rgb).tobytes()).decode()) | |
| page = VIEWER.replace("__DATA__", json.dumps(data)) | |
| return (f'<iframe srcdoc="{html.escape(page, quote=True)}" ' | |
| 'style="width:100%;height:560px;border:0;border-radius:10px;background:#0b0d10" allowfullscreen></iframe>') | |
| EMPTY_VIEWER = ('<div style="height:560px;display:flex;align-items:center;justify-content:center;' | |
| 'border-radius:10px;background:#0b0d10;color:#8a9099;font:14px system-ui">' | |
| 'The reconstructed point cloud and cameras will appear here.</div>') | |
| def colorize_depth(depth, conf): | |
| d = depth.copy() | |
| valid = np.isfinite(d) & (d > 0) | |
| if not valid.any(): | |
| return np.zeros((*d.shape, 3), np.uint8) | |
| lo, hi = np.percentile(d[valid], 2), np.percentile(d[valid], 98) | |
| n = np.clip((d - lo) / max(hi - lo, 1e-9), 0, 1) | |
| img = cv2.applyColorMap(((1 - n) * 255).astype(np.uint8), cv2.COLORMAP_TURBO) | |
| return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| # ---------------------------------------------------------------- inputs | |
| def _collect_paths(files): | |
| if not files: | |
| raise gr.Error("Upload at least 2 images of the same scene (or one video).") | |
| paths = [f if isinstance(f, str) else f.name for f in files] | |
| videos = [p for p in paths if p.lower().endswith(VIDEO_EXTENSIONS)] | |
| if videos: | |
| frames = list_images(videos[0]) # every frame of the first video | |
| stride = max(1, int(np.ceil(len(frames) / 32))) | |
| return frames[::stride][:32] | |
| imgs = [p for p in paths if p.lower().endswith((".jpg", ".jpeg", ".png", ".webp", ".bmp"))] | |
| if len(imgs) < 2: | |
| raise gr.Error("Solv3R needs at least 2 images of the same scene.") | |
| if len(imgs) > MAX_VIEWS: | |
| idx = np.linspace(0, len(imgs) - 1, MAX_VIEWS).round().astype(int) | |
| imgs = [imgs[i] for i in idx] | |
| gr.Warning(f"Using {MAX_VIEWS} evenly spaced images of the {len(paths)} uploaded.") | |
| return imgs | |
| def _duration(files, calls=16, adaptive=False, keep=0.5, sky=False, *args, **kwargs): | |
| n = min(len(files) if files else 2, MAX_VIEWS) | |
| calls = 64 if adaptive else int(calls) | |
| # measured on ZeroGPU: 39 views / 16 calls -> 2.9 s warm, 13.4 s cold with sky masking | |
| return int(min(90, 8 + n * 0.12 + calls * n * 0.004 + (n * 0.15 if sky else 0))) | |
| def reconstruct(files: list, calls: int = 16, adaptive: bool = False, keep: float = 0.5, sky: bool = False, | |
| progress=gr.Progress()): | |
| """Reconstruct a 3D scene (point cloud, depth maps, camera poses) from multiple views with Solv3R. | |
| Args: | |
| files: two or more images of the same scene (or one video, sampled to at most 32 frames). | |
| calls: number of calls of the looped solver block (16 = the paper's default). | |
| adaptive: stop early once one call changes the state by less than 2% (at most 64 calls). | |
| keep: fraction of the most confident pixels per view kept in the point cloud. | |
| sky: remove the sky from the point cloud (outdoor captures). | |
| """ | |
| t0 = time.perf_counter() | |
| paths = _collect_paths(files) | |
| images = load_images(paths) | |
| progress(0.1, desc=f"Encoding {len(paths)} views") | |
| ctx = model.encode(images) | |
| z = model.start(ctx) | |
| n_calls = 64 if adaptive else int(calls) | |
| change, call = float("nan"), 0 | |
| with autocast(): | |
| for call in range(1, n_calls + 1): | |
| z_next = model.step(z, ctx) | |
| change = float((z_next.float() - z.float()).norm() / z.float().norm()) | |
| z = z_next | |
| progress(0.1 + 0.7 * call / n_calls, desc=f"Solver call {call}") | |
| if adaptive and call >= 4 and change < 0.02: | |
| break | |
| pred = model.decode(z, ctx) | |
| torch.cuda.synchronize() | |
| t_model = time.perf_counter() - t0 | |
| hw = (ctx.H, ctx.W) | |
| c2w, K = cameras_from_pose_enc(pred["pose_enc"], hw) | |
| c2w, K = c2w.numpy(), K.numpy() | |
| depth = pred["depth"][..., 0].float().cpu().numpy() | |
| conf = pred["depth_conf"].float().cpu().numpy() | |
| points = pred["world_points"].float().cpu().numpy() | |
| del pred, z, ctx | |
| progress(0.85, desc="Building point cloud") | |
| valid = ~sky_mask(images) if sky else None | |
| xyz, rgb = confident_points(points, conf, images, float(keep), valid) | |
| out = tempfile.mkdtemp(prefix="solv3r_") | |
| ply = os.path.join(out, "solv3r_pointcloud.ply") | |
| write_ply(ply, xyz, rgb) | |
| npz = os.path.join(out, "solv3r_cameras_depth.npz") | |
| np.savez_compressed(npz, c2w=c2w, K=K, depth=depth.astype(np.float16), depth_conf=conf.astype(np.float16), | |
| names=np.array([os.path.basename(p) for p in paths])) | |
| rgb_imgs = (images.permute(0, 2, 3, 1).clamp(0, 1).numpy() * 255).astype(np.uint8) | |
| gallery = [] | |
| for i in range(len(paths)): | |
| gallery.append((np.concatenate([rgb_imgs[i], colorize_depth(depth[i], conf[i])], axis=1), | |
| f"view {i}: {os.path.basename(paths[i])}")) | |
| cams = [] | |
| for i in range(len(paths)): | |
| t = c2w[i, :3, 3] | |
| cams.append([i, os.path.basename(paths[i]), round(float(t[0]), 4), round(float(t[1]), 4), round(float(t[2]), 4), | |
| round(float(K[i, 0, 0]), 1), round(float(K[i, 1, 1]), 1)]) | |
| viewer = build_viewer(xyz, rgb, c2w, K, hw) | |
| status = (f"{len(paths)} views at {hw[1]}x{hw[0]} · {call} solver calls (last relative change {change:.4f}) · " | |
| f"{len(xyz):,} points · solve {t_model:.1f}s, total {time.perf_counter() - t0:.1f}s") | |
| print(status) | |
| return viewer, gallery, cams, [ply, npz], status | |
| # ---------------------------------------------------------------- UI | |
| def _example(folder): | |
| return sorted(glob.glob(os.path.join("examples", folder, "*.jpg"))) | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# Solv3R: 3D Looped Transformer as a Multi-view Solver\n" | |
| "Upload several photos of a scene (or a short video) to recover a colored point cloud, per-view depth " | |
| "and camera poses. " | |
| "[Paper](https://arxiv.org/abs/2611.02857) · [Project page](https://ubc-vision.github.io/solv3r/) · " | |
| "[Code](https://github.com/ubc-vision/solv3r) · [Model](https://huggingface.co/jgkwak/solv3r)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| files = gr.File(label="Images (2–64) or one video", file_count="multiple", | |
| file_types=["image", "video"], height=260) | |
| run = gr.Button("Reconstruct", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| calls = gr.Slider(1, 48, value=16, step=1, label="Solver calls", | |
| info="Calls of the looped block (paper default: 16).") | |
| adaptive = gr.Checkbox(False, label="Adaptive stopping", | |
| info="Stop once a call changes the state by < 2% (max 64 calls).") | |
| keep = gr.Slider(0.1, 1.0, value=0.5, step=0.05, label="Confidence keep ratio", | |
| info="Fraction of the most confident pixels per view kept in the cloud.") | |
| sky = gr.Checkbox(False, label="Remove sky", info="For outdoor captures.") | |
| status = gr.Textbox(label="Status", interactive=False) | |
| with gr.Column(scale=2): | |
| viewer = gr.HTML(EMPTY_VIEWER) | |
| with gr.Tabs(): | |
| with gr.Tab("Depth maps"): | |
| gallery = gr.Gallery(label="Input | predicted depth", columns=3, height="auto") | |
| with gr.Tab("Cameras"): | |
| cams = gr.Dataframe(headers=["view", "image", "tx", "ty", "tz", "fx", "fy"], | |
| label="Camera centers (world frame = first view) and focal lengths (px)") | |
| with gr.Tab("Downloads"): | |
| downloads = gr.File(label="Point cloud (.ply) and cameras + depth (.npz)", file_count="multiple") | |
| outputs = [viewer, gallery, cams, downloads, status] | |
| run.click(reconstruct, inputs=[files, calls, adaptive, keep, sky], outputs=outputs, api_name="reconstruct") | |
| gr.Examples( | |
| examples=[[_example("park"), 16, False, 0.5, True], [_example("desk"), 16, False, 0.5, False]], | |
| inputs=[files, calls, adaptive, keep, sky], | |
| outputs=outputs, | |
| fn=reconstruct, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples (39 drone photos of a park · 24 frames of a desk video, from the Solv3R repo)", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |