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 = """
Solv3R
"""
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("')
EMPTY_VIEWER = ('
'
'The reconstructed point cloud and cameras will appear here.
')
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)))
@spaces.GPU(duration=_duration)
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)