Pointmap: byte-accurate MoGe-2 mesh recipe

Match upstream utils3d.numpy exactly:
- _depth_edge: NaN-pad + nanmax (boundary pixels excluded), formula
diff / depth > rtol (raw depth in denominator).
- Quad winding [TL, BL, BR, TR] split as fan from TL into triangles
[TL, BL, BR] and [TL, BR, TR].

Drop dead helpers (_camera_marker, _triangulate_grid, _glb_inject_unlit).

app.py

@@ -170,87 +170,46 @@ def _depth_to_rgb(depth: np.ndarray, mask: np.ndarray) -> np.ndarray:
 
 
 # -----------------------------------------------------------------------------
-# Point cloud export — trimesh → .glb (much faster than Open3D .ply for Three.js)
-
-def _camera_marker(radius: float = 0.025, n_points: int = 600,
-                   color=(51, 140, 245)):
-    """Tiny slate-blue Fibonacci sphere marking the camera. Returns (verts, cols)."""
-    i = np.arange(n_points)
-    phi = np.arccos(1 - 2 * (i + 0.5) / n_points)
-    theta = np.pi * (1 + 5 ** 0.5) * (i + 0.5)
-    verts = np.stack([
-        radius * np.sin(phi) * np.cos(theta),
-        radius * np.sin(phi) * np.sin(theta),
-        radius * np.cos(phi),
-    ], axis=1).astype(np.float32)
-    cols = np.tile(np.array(color + (255,), dtype=np.uint8), (n_points, 1))
-    return verts, cols
-
-
-def _triangulate_grid(pointmap_hwc: np.ndarray, mask_hw: np.ndarray,
-                      max_edge: float = 0.04):
-    """Build a triangulated mesh from the (H, W) pointmap grid.
-
-    Returns (verts, uvs, faces, vertex_normals). Each valid pixel → vertex;
-    adjacent valid pixels form quads (2 tris). Long-edge triangles are dropped
-    to kill stretched skin at depth jumps. Vertex normals are computed from the
-    pointmap's spatial gradient (smooth, per-pixel) instead of inferred from
-    triangle face normals (which would give flat-shaded facets).
+# Mesh export — MoGe-2's recipe (trimesh → .glb)
+#
+# We build a regular grid mesh from the (H, W) pointmap: each valid pixel is a
+# vertex, adjacent valid pixels form quads → 2 triangles each. The trick that
+# makes MoGe-2's meshes look clean (no stretched-skin facets at depth jumps,
+# no ragged silhouette) is what they call `mask_cleaned`:
+#
+#     mask_cleaned = mask & ~depth_edge(depth, rtol=0.04)
+#
+# i.e. drop pixels sitting on a depth discontinuity *before* triangulation, so
+# no triangle ever spans one. We don't post-filter triangles by edge length.
+
+def _depth_edge(depth: np.ndarray, rtol: float = 0.04, kernel_size: int = 3) -> np.ndarray:
+    """NumPy port of `utils3d.numpy.depth_edge` (rtol-only).
+
+    For each pixel, look at the kernel×kernel window around it; if
+    (max − min)/depth > rtol, mark it as a depth-edge pixel.
     """
-    H, W = pointmap_hwc.shape[:2]
-    z = pointmap_hwc[:, :, 2]
-    valid = mask_hw & np.isfinite(pointmap_hwc).all(axis=2) & (z > 0.05) & (z < 25.0)
-
-    # Smooth per-pixel normals from cross product of x- and y- spatial gradients.
-    px = np.zeros_like(pointmap_hwc, dtype=np.float32)
-    py = np.zeros_like(pointmap_hwc, dtype=np.float32)
-    px[:, 1:-1] = (pointmap_hwc[:, 2:] - pointmap_hwc[:, :-2]) * 0.5
-    py[1:-1, :] = (pointmap_hwc[2:, :] - pointmap_hwc[:-2, :]) * 0.5
-    n_grid = np.cross(px, py)
-    n_grid /= np.linalg.norm(n_grid, axis=2, keepdims=True).clip(min=1e-8)
-
-    idx_map = np.full((H, W), -1, dtype=np.int64)
-    yy, xx = np.where(valid)
-    idx_map[yy, xx] = np.arange(len(yy))
-
-    verts = pointmap_hwc[yy, xx].astype(np.float32)              # (N, 3)
-    normals = n_grid[yy, xx].astype(np.float32)                   # (N, 3)
-    uvs = np.stack([xx / max(W - 1, 1), yy / max(H - 1, 1)],
-                   axis=1).astype(np.float32)                    # (N, 2)
-
-    a = idx_map[:-1, :-1]; b = idx_map[:-1, 1:]
-    c = idx_map[1:, :-1];  d = idx_map[1:, 1:]
-    quad_valid = (a != -1) & (b != -1) & (c != -1) & (d != -1)
-    a_v, b_v, c_v, d_v = a[quad_valid], b[quad_valid], c[quad_valid], d[quad_valid]
-    tri1 = np.stack([a_v, c_v, b_v], axis=1)
-    tri2 = np.stack([b_v, c_v, d_v], axis=1)
-    faces = np.concatenate([tri1, tri2], axis=0)
-
-    p0 = verts[faces[:, 0]]; p1 = verts[faces[:, 1]]; p2 = verts[faces[:, 2]]
-    e01 = np.linalg.norm(p1 - p0, axis=1)
-    e12 = np.linalg.norm(p2 - p1, axis=1)
-    e20 = np.linalg.norm(p0 - p2, axis=1)
-    keep = (e01 < max_edge) & (e12 < max_edge) & (e20 < max_edge)
-    faces = faces[keep].astype(np.int64)
-    return verts, uvs, faces, normals
+    pad = kernel_size // 2
+    # NaN-pad + nanmax/nanmin = ignore out-of-bounds pixels at image borders
+    # (matches upstream `utils3d.numpy.max_pool_1d`).
+    padded = np.pad(depth.astype(np.float32), pad, mode="constant", constant_values=np.nan)
+    windows = np.lib.stride_tricks.sliding_window_view(padded, (kernel_size, kernel_size))
+    d_max = np.nanmax(windows, axis=(-2, -1))
+    d_min = np.nanmin(windows, axis=(-2, -1))
+    with np.errstate(divide="ignore", invalid="ignore"):
+        rel = (d_max - d_min) / depth
+    return np.nan_to_num(rel, nan=0.0, posinf=0.0, neginf=0.0) > rtol
 
 
 def _make_glb(image_pil_texture: Image.Image, pointmap_hwc: np.ndarray,
-              mask_hw: np.ndarray, max_edge: float = 0.04) -> str:
-    """Triangulated mesh, UV-textured with the input image — MoGe-2's recipe.
-
-    Each valid pixel = vertex; adjacent valid pixels form quads → 2 triangles.
-    Triangles whose edges exceed `max_edge` (meters) are dropped to kill
-    stretched skin at depth jumps. The input image is used as the GLB's albedo
-    texture (per-triangle PBR sampling), and trimesh's lazy vertex_normals get
-    exported so Three.js applies smooth shading instead of flat facets.
-    """
+              mask_hw: np.ndarray, rtol: float = 0.04) -> str:
+    """Build a UV-textured triangulated mesh and export to .glb (MoGe-2 recipe)."""
     H, W = pointmap_hwc.shape[:2]
     image_native = image_pil_texture.resize((W, H), Image.LANCZOS)
 
-    # Triangulate the (H, W) grid over the foreground mask.
     z = pointmap_hwc[:, :, 2]
     valid = mask_hw & np.isfinite(pointmap_hwc).all(axis=2) & (z > 0.05) & (z < 25.0)
+    valid &= ~_depth_edge(z, rtol=rtol)
+
     idx_map = np.full((H, W), -1, dtype=np.int64)
     yy, xx = np.where(valid)
     idx_map[yy, xx] = np.arange(len(yy))
@@ -258,30 +217,22 @@ def _make_glb(image_pil_texture: Image.Image, pointmap_hwc: np.ndarray,
     verts = pointmap_hwc[yy, xx].astype(np.float32)
     uvs = np.stack([xx / max(W - 1, 1), yy / max(H - 1, 1)], axis=1).astype(np.float32)
 
-    a = idx_map[:-1, :-1]; b = idx_map[:-1, 1:]
-    c = idx_map[1:, :-1];  d = idx_map[1:, 1:]
-    quad_valid = (a != -1) & (b != -1) & (c != -1) & (d != -1)
-    a_v, b_v, c_v, d_v = a[quad_valid], b[quad_valid], c[quad_valid], d[quad_valid]
-    tri1 = np.stack([a_v, c_v, b_v], axis=1)
-    tri2 = np.stack([b_v, c_v, d_v], axis=1)
-    faces = np.concatenate([tri1, tri2], axis=0)
-
-    p0 = verts[faces[:, 0]]; p1 = verts[faces[:, 1]]; p2 = verts[faces[:, 2]]
-    e01 = np.linalg.norm(p1 - p0, axis=1)
-    e12 = np.linalg.norm(p2 - p1, axis=1)
-    e20 = np.linalg.norm(p0 - p2, axis=1)
-    keep = (e01 < max_edge) & (e12 < max_edge) & (e20 < max_edge)
-    faces = faces[keep].astype(np.int64)
+    # Quad order matches upstream `utils3d.numpy.image_mesh`: [TL, BL, BR, TR],
+    # split into triangles as fan from TL → [TL, BL, BR] and [TL, BR, TR].
+    tl = idx_map[:-1, :-1]; tr = idx_map[:-1, 1:]
+    bl = idx_map[1:, :-1];  br = idx_map[1:, 1:]
+    quad_valid = (tl != -1) & (tr != -1) & (bl != -1) & (br != -1)
+    tl_v, tr_v, bl_v, br_v = tl[quad_valid], tr[quad_valid], bl[quad_valid], br[quad_valid]
+    tri1 = np.stack([tl_v, bl_v, br_v], axis=1)
+    tri2 = np.stack([tl_v, br_v, tr_v], axis=1)
+    faces = np.concatenate([tri1, tri2], axis=0).astype(np.int64)
 
     # MoGe-2: y/z flip on positions, v-flip on UVs.
-    flip = np.array([1.0, -1.0, -1.0], dtype=np.float32)
-    verts = verts * flip
+    verts = verts * np.array([1.0, -1.0, -1.0], dtype=np.float32)
     centroid = verts.mean(axis=0).astype(np.float32) if len(verts) else np.zeros(3, np.float32)
     verts = verts - centroid
-
     uvs = uvs * np.array([1.0, -1.0], dtype=np.float32) + np.array([0.0, 1.0], dtype=np.float32)
 
-    # PBR with image as albedo texture — MoGe's exact material settings.
     material = trimesh.visual.material.PBRMaterial(
         baseColorTexture=image_native,
         metallicFactor=0.5,
@@ -291,65 +242,18 @@ def _make_glb(image_pil_texture: Image.Image, pointmap_hwc: np.ndarray,
     visual = trimesh.visual.texture.TextureVisuals(uv=uvs, material=material)
 
     mesh = trimesh.Trimesh(
-        vertices=verts, faces=faces, visual=visual, process=False,
+        vertices=verts,
+        faces=faces,
+        vertex_normals=None,  # MoGe-2 leaves this to the GLB consumer
+        visual=visual,
+        process=False,
     )
-    # Touch vertex_normals so trimesh caches them; the GLB exporter reads
-    # the cached normals and writes them into the file → smooth shading
-    # in Three.js (instead of the per-face fallback that looked like facets).
-    _ = mesh.vertex_normals
 
     out_path = tempfile.NamedTemporaryFile(delete=False, suffix=".glb").name
     mesh.export(out_path)
     return out_path
 
 
-# -----------------------------------------------------------------------------
-# GLB post-processing: inject KHR_materials_unlit extension so Three.js skips
-# all lighting calculations and renders surfels as their raw vertex colour.
-
-import struct
-import json
-
-
-def _glb_inject_unlit(glb_path: str) -> None:
-    """Patch a binary GLB to mark every material as KHR_materials_unlit.
-    glTF spec for the extension: surfaces render with `baseColor` only,
-    no shading from light sources or normals."""
-    with open(glb_path, "rb") as f:
-        data = f.read()
-
-    # GLB header: magic | version | total_length  (12 bytes)
-    if data[:4] != b"glTF":
-        return
-    # First chunk: length(4) | type(4) | payload — JSON chunk.
-    json_len, json_type = struct.unpack_from("<II", data, 12)
-    if json_type != 0x4E4F534A:  # "JSON"
-        return
-    json_bytes = data[20 : 20 + json_len]
-    bin_tail = data[20 + json_len :]   # everything after = BIN chunk(s)
-
-    gltf = json.loads(json_bytes.rstrip(b" \x00").decode("utf-8"))
-    used = gltf.setdefault("extensionsUsed", [])
-    if "KHR_materials_unlit" not in used:
-        used.append("KHR_materials_unlit")
-    for mat in gltf.get("materials", []):
-        mat.setdefault("extensions", {})["KHR_materials_unlit"] = {}
-
-    new_json = json.dumps(gltf, separators=(",", ":")).encode("utf-8")
-    pad = (4 - len(new_json) % 4) % 4
-    new_json += b" " * pad
-
-    new_total = 12 + 8 + len(new_json) + len(bin_tail)
-    out = (
-        b"glTF" + struct.pack("<II", 2, new_total)
-        + struct.pack("<II", len(new_json), 0x4E4F534A)
-        + new_json
-        + bin_tail
-    )
-    with open(glb_path, "wb") as f:
-        f.write(out)
-
-
 # -----------------------------------------------------------------------------
 # Gradio handler