Pointmap: depth heatmap on the right with solid grey bg; .ply still in accordion

app.py

@@ -134,6 +134,26 @@ def _foreground_mask(image_pil: Image.Image, target_h: int, target_w: int) -> np
     return (out.argmax(dim=1)[0] > 0).cpu().numpy()
 
 
+def _depth_to_rgb(depth: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    """Inverse-depth turbo colormap (matches sapiens2 vis_pointmap.py).
+    Background pixels are left at 0 — caller should overlay them."""
+    valid = np.isfinite(depth) & (depth > 1e-3) & mask
+    rgb = np.zeros((*depth.shape, 3), dtype=np.uint8)
+    if not valid.any():
+        return rgb
+    inv = np.zeros_like(depth, dtype=np.float32)
+    inv[valid] = 1.0 / depth[valid]
+    p1, p99 = np.percentile(inv[valid], [1, 99])
+    lo, hi = float(p1), float(p99)
+    if hi <= lo:
+        hi = lo + 1e-3
+    norm = ((inv - lo) / (hi - lo)).clip(0, 1)
+    grey = (norm * 255.0).astype(np.uint8)
+    color = cv2.applyColorMap(grey, cv2.COLORMAP_TURBO)[:, :, ::-1]  # cv2 is BGR → RGB
+    rgb[valid] = color[valid]
+    return rgb
+
+
 # -----------------------------------------------------------------------------
 # Point cloud export (camera marker + cloud, native-res grid)
 
@@ -194,13 +214,24 @@ def predict(image: Image.Image, size: str):
     image_bgr = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
 
     model = _get_pointmap_model(size)
-    pointmap = _estimate_pointmap(image_bgr, model)        # (H_n, W_n, 3) — at most 1024 in either dim
+    pointmap = _estimate_pointmap(image_bgr, model)        # (H_n, W_n, 3) at most 1024 in either dim
     h_n, w_n = pointmap.shape[:2]
 
     mask = _foreground_mask(image_pil, h_n, w_n)            # native-res mask, fast
-    ply_path = _make_ply(image_pil, pointmap, mask)         # native-res .ply
 
-    return ply_path, ply_path
+    # Depth heatmap (right pane). Solid mid-grey background with the foreground
+    # turbo-coloured by inverse depth. Mirrors sapiens2 vis_pointmap.py colormap.
+    depth = pointmap[:, :, 2]
+    depth_rgb = _depth_to_rgb(depth, mask)
+    BG_GREY = 200
+    depth_rgb[~mask] = BG_GREY
+    w0, h0 = image_pil.size
+    depth_pil = Image.fromarray(depth_rgb).resize((w0, h0), Image.LANCZOS)
+
+    # PLY (download in accordion). Native-res, ≤200K points.
+    ply_path = _make_ply(image_pil, pointmap, mask)
+
+    return depth_pil, ply_path
 
 
 # -----------------------------------------------------------------------------
@@ -268,14 +299,7 @@ with gr.Blocks(title="Sapiens2 Pointmap", theme=gr.themes.Soft(), css=CUSTOM_CSS
 
     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],
-            display_mode="point_cloud",
-            zoom_speed=0.7,
-            pan_speed=0.5,
-        )
+        out_img = gr.Image(label="Depth (Z)", type="pil", height=640)
 
     with gr.Row():
         size = gr.Radio(
@@ -291,7 +315,7 @@ with gr.Blocks(title="Sapiens2 Pointmap", theme=gr.themes.Soft(), css=CUSTOM_CSS
     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])
+    run.click(predict, inputs=[inp, size], outputs=[out_img, out_ply_file])
 
 
 if __name__ == "__main__":