Add PXDepth demo

app.py

@@ -39,10 +39,41 @@ OUTPUT_MAX_AGE = 60 * 60
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 CSS = """
+html, body {
+    height: auto !important;
+    min-height: 100% !important;
+    overflow-y: auto !important;
+    overscroll-behavior-y: auto !important;
+}
+.gradio-container {
+    height: auto !important;
+    min-height: 100vh !important;
+    overflow: visible !important;
+}
 #pxdepth-demo { max-width: 1280px; margin: 0 auto; }
 #img-display-input, #img-display-output { max-height: 72vh; }
 #img-display-output img { object-fit: contain !important; }
-#model-3d { min-height: 60vh; }
+#model-3d { min-height: 55vh; }
+"""
+
+PAGE_JS = """
+() => {
+    document.documentElement.style.overflowY = "auto";
+    document.body.style.overflowY = "auto";
+    const install = () => {
+        const viewer = document.querySelector("#model-3d");
+        if (!viewer || viewer.dataset.pageWheel === "true") return;
+        viewer.dataset.pageWheel = "true";
+        viewer.addEventListener("wheel", (event) => {
+            if (event.ctrlKey || event.metaKey) return;
+            event.preventDefault();
+            event.stopImmediatePropagation();
+            window.scrollBy({ top: event.deltaY, left: 0, behavior: "auto" });
+        }, { passive: false, capture: true });
+    };
+    install();
+    new MutationObserver(install).observe(document.body, { childList: true, subtree: true });
+}
 """
 
 
@@ -125,6 +156,103 @@ def sample_points(
     return points[indices], colors[indices]
 
 
+def filter_flying_points(
+    points: np.ndarray,
+    colors: np.ndarray,
+    neighbors: int = 30,
+    std_ratio: float = 2.0,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Remove sparse statistical outliers from an already sampled cloud."""
+    if points.shape[0] <= neighbors + 1:
+        return points, colors
+
+    from scipy.spatial import cKDTree
+
+    tree = cKDTree(points.astype(np.float64, copy=False))
+    mean_distance = np.empty(points.shape[0], dtype=np.float32)
+    for start in range(0, points.shape[0], 100_000):
+        stop = min(start + 100_000, points.shape[0])
+        try:
+            distances, _ = tree.query(
+                points[start:stop],
+                k=neighbors + 1,
+                workers=-1,
+            )
+        except TypeError:
+            distances, _ = tree.query(points[start:stop], k=neighbors + 1)
+        mean_distance[start:stop] = np.asarray(
+            distances[:, 1:],
+            dtype=np.float32,
+        ).mean(axis=1)
+    finite = np.isfinite(mean_distance)
+    if not finite.any():
+        return points, colors
+    values = mean_distance[finite]
+    threshold = float(values.mean() + std_ratio * values.std())
+    keep = finite & (mean_distance <= threshold)
+    return (points[keep], colors[keep]) if keep.any() else (points, colors)
+
+
+def write_splat(
+    path: Path,
+    points: np.ndarray,
+    colors: np.ndarray,
+    point_size: float,
+) -> None:
+    """Write a compact Gaussian-splat viewer file with adjustable point size."""
+    display_points = points * np.array([1.0, -1.0, -1.0], dtype=np.float32)
+    low, high = np.percentile(display_points, [1.0, 99.0], axis=0)
+    extent = max(float(np.max(high - low)), 1e-6)
+    spacing = extent / max(float(np.sqrt(points.shape[0])), 1.0)
+    scale = max(spacing * float(point_size) * 0.65, 1e-7)
+
+    dtype = np.dtype(
+        [
+            ("position", "<f4", (3,)),
+            ("scale", "<f4", (3,)),
+            ("color", "u1", (4,)),
+            ("rotation", "u1", (4,)),
+        ],
+        align=False,
+    )
+    splats = np.empty(points.shape[0], dtype=dtype)
+    splats["position"] = display_points
+    splats["scale"] = scale
+    splats["color"][:, :3] = colors
+    splats["color"][:, 3] = 255
+    splats["rotation"] = np.array([255, 128, 128, 128], dtype=np.uint8)
+    splats.tofile(path)
+
+
+def update_viewer(
+    cache_path: Optional[str],
+    filter_points: bool,
+    point_size: float,
+    max_points: int,
+) -> Optional[str]:
+    """Rebuild the viewer from cached points without running either model."""
+    if not cache_path or not Path(cache_path).is_file():
+        return None
+
+    with np.load(cache_path) as cache:
+        points = cache["points"]
+        colors = cache["colors"]
+    points, colors = sample_points(points, colors, int(max_points))
+    if filter_points:
+        points, colors = filter_flying_points(points, colors)
+    if points.shape[0] == 0:
+        raise gr.Error("No points remain after filtering.")
+
+    cache_file = Path(cache_path)
+    tag = f"{int(max_points)}_{int(filter_points)}_{float(point_size):.2f}"
+    viewer_path = cache_file.with_name(f"pointcloud_viewer_{tag}.splat")
+    write_splat(viewer_path, points, colors, point_size)
+    for old_path in cache_file.parent.glob("pointcloud_viewer_*.splat"):
+        if old_path != viewer_path:
+            old_path.unlink(missing_ok=True)
+    return str(viewer_path)
+
+
 @gpu
 @torch.inference_mode()
 def predict_gpu(image: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
@@ -163,8 +291,10 @@ def predict_gpu(image: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
 
 def on_submit(
     image: Optional[np.ndarray],
-    max_points: int,
     apply_mask: bool,
+    filter_points: bool,
+    point_size: float,
+    max_points: int,
     request: gr.Request,
 ):
     """Run inference, build visualizations, and export downloadable files."""
@@ -199,11 +329,11 @@ def on_submit(
     depth_vis = colorize_depth(np.where(mask_np, depth_np, np.inf), mask=None)
 
     output = session_dir(request)
-    depth_npy = output / "metric_depth.npy"
+    depth_npy = output / "depth.npy"
     depth_png = output / "depth.png"
     mask_png = output / "mask.png"
     ply_path = output / "pointcloud.ply"
-    glb_path = output / "pointcloud_viewer.glb"
+    cache_path = output / "viewer_data.npz"
 
     np.save(depth_npy, depth_np)
     Image.fromarray(depth_vis).save(depth_png)
@@ -216,20 +346,17 @@ def on_submit(
     if points_full.shape[0] == 0:
         raise gr.Error("No valid 3D points were produced for this image.")
     write_point_cloud_ply(ply_path, points_full, colors_full)
-
-    import trimesh
-
-    viewer_points, viewer_colors = sample_points(
-        points_full,
-        colors_full,
-        int(max_points),
+    colors_uint8 = np.clip(colors_full * 255.0, 0, 255).astype(np.uint8)
+    np.savez(cache_path, points=points_full.astype(np.float32), colors=colors_uint8)
+    viewer_path = update_viewer(
+        str(cache_path),
+        filter_points,
+        point_size,
+        max_points,
     )
-    viewer_points = viewer_points * np.array([1.0, -1.0, -1.0], np.float32)
-    viewer_colors = np.clip(viewer_colors * 255.0, 0, 255).astype(np.uint8)
-    trimesh.PointCloud(viewer_points, colors=viewer_colors).export(glb_path)
 
     files = [str(depth_png), str(depth_npy), str(mask_png), str(ply_path)]
-    return (image, depth_vis), str(glb_path), files
+    return (image, depth_vis), viewer_path, files, str(cache_path)
 
 
 def build_demo() -> gr.Blocks:
@@ -240,7 +367,8 @@ See the [paper](https://arxiv.org/abs/2608.16984),
 [project page](https://yuanzhy29.github.io/PXDepth-Page/), and
 [GitHub repository](https://github.com/yuanzhy29/PXDepth).
 """
-    with gr.Blocks(theme=gr.themes.Soft(), css=CSS) as demo:
+    with gr.Blocks(theme=gr.themes.Soft(), css=CSS, js=PAGE_JS) as demo:
+        viewer_cache = gr.State(value=None)
         with gr.Column(elem_id="pxdepth-demo"):
             gr.Markdown("# PXDepth")
             gr.Markdown(description)
@@ -255,17 +383,30 @@ See the [paper](https://arxiv.org/abs/2608.16984),
                         elem_id="img-display-input",
                     )
                     with gr.Accordion(label="Settings", open=False):
+                        apply_mask = gr.Checkbox(
+                            label="Apply valid-depth mask to point cloud",
+                            value=True,
+                        )
+                        filter_points = gr.Checkbox(
+                            label="Filter Flying Points",
+                            info="Statistical outlier filtering; does not rerun the model.",
+                            value=False,
+                        )
+                        point_size = gr.Slider(
+                            0.25,
+                            3.0,
+                            value=1.0,
+                            step=0.05,
+                            label="Point Size",
+                            info="Updates only the 3D viewer.",
+                        )
                         max_points = gr.Slider(
                             50_000,
                             500_000,
                             value=200_000,
                             step=50_000,
                             label="3D Viewer Max Points",
-                            info="The downloaded PLY retains all valid points.",
-                        )
-                        apply_mask = gr.Checkbox(
-                            label="Apply valid-depth mask to point cloud",
-                            value=True,
+                            info="Updates only the viewer; the downloaded PLY retains all valid points.",
                         )
                     submit = gr.Button("Predict", variant="primary")
 
@@ -275,7 +416,7 @@ See the [paper](https://arxiv.org/abs/2608.16984),
                             model_3d = gr.Model3D(
                                 label="3D Point Map",
                                 clear_color=(1.0, 1.0, 1.0, 1.0),
-                                height="60vh",
+                                height="55vh",
                                 elem_id="model-3d",
                             )
                         with gr.Tab("Depth"):
@@ -308,11 +449,30 @@ See the [paper](https://arxiv.org/abs/2608.16984),
 
             submit.click(
                 on_submit,
-                [input_image, max_points, apply_mask],
-                [depth_map, model_3d, downloads],
+                [input_image, apply_mask, filter_points, point_size, max_points],
+                [depth_map, model_3d, downloads, viewer_cache],
                 show_progress="full",
                 concurrency_limit=1,
             )
+            viewer_inputs = [viewer_cache, filter_points, point_size, max_points]
+            filter_points.change(
+                update_viewer,
+                viewer_inputs,
+                model_3d,
+                show_progress="minimal",
+            )
+            point_size.release(
+                update_viewer,
+                viewer_inputs,
+                model_3d,
+                show_progress="minimal",
+            )
+            max_points.release(
+                update_viewer,
+                viewer_inputs,
+                model_3d,
+                show_progress="minimal",
+            )
     return demo
 
 

requirements.txt

@@ -8,6 +8,5 @@ huggingface_hub>=0.28
 matplotlib>=3.8
 pillow>=10.0
 scipy>=1.12
-trimesh>=4.5
 git+https://github.com/EasternJournalist/utils3d.git@3fab839f0be9931dac7c8488eb0e1600c236e183
 git+https://github.com/microsoft/MoGe.git@42acd8f46e974f5e0548ecd72315d1d7df2cb6f4