Upload 2 files

app.py

@@ -6,21 +6,16 @@ import gradio as gr
 import numpy as np
 import torch
 import cv2
+import open3d as o3d
+import trimesh
 
 from moge.model.v2 import MoGeModel
 
-
-# ---------- Model setup ----------
-
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 @torch.no_grad()
 def load_model() -> MoGeModel:
-    """
-    Load MoGe from the HF repo 'Ruicheng/moge-2-vitl-normal'.
-    This will download model.pt the first time and cache it.
-    """
     print(f"Loading MoGe model on device: {DEVICE}")
     model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal")
     model = model.to(DEVICE)
@@ -31,8 +26,6 @@ def load_model() -> MoGeModel:
 MODEL = load_model()
 
 
-# ---------- Helper: run MoGe & get point cloud ----------
-
 @torch.no_grad()
 def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
     """
@@ -42,37 +35,31 @@ def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
         points: (N, 3) float32 XYZ
         colors: (N, 3) uint8 RGB
     """
-
-    # Convert to float tensor [0, 1], CHW, batch
     img = image.astype(np.float32) / 255.0
-    tensor = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).to(DEVICE)  # (1,3,H,W)
+    tensor = (
+        torch.from_numpy(img)
+        .permute(2, 0, 1)
+        .unsqueeze(0)
+        .to(DEVICE)  # (1,3,H,W)
+    )
 
-    # --- Run MoGe ---
     out = MODEL.infer(tensor)
 
-    # --- DEBUG: log what MoGe actually returned ---
     print("MoGe output keys:", list(out.keys()))
-    shaped = {}
-    for k, v in out.items():
-        if torch.is_tensor(v):
-            shaped[k] = (v.shape, v.dtype, float(v.min()), float(v.max()))
-        else:
-            shaped[k] = type(v).__name__
-    print("MoGe output summary:", shaped)
 
-    # --- Try several common patterns ---
+    # You already have this part working;
+    # keep your existing logic if it's different.
+    # Here’s a generic version that assumes out["pcd"] (B,N,6) or out["points"]/out["colors"].
 
     points = None
     colors = None
 
-    # 1) Single tensor with xyzrgb in last dim: (B, N, 6)
     if "pcd" in out:
         pcd = out["pcd"]
         if pcd.ndim == 3 and pcd.shape[-1] >= 3:
-            # remove batch
             if pcd.shape[0] == 1:
                 pcd = pcd[0]
-            pcd_np = pcd.detach().cpu().float().numpy()  # (N, C)
+            pcd_np = pcd.detach().cpu().float().numpy()
             points = pcd_np[:, :3]
             if pcd_np.shape[1] >= 6:
                 cols = pcd_np[:, 3:6]
@@ -80,7 +67,6 @@ def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
                     cols = (cols * 255.0).clip(0, 255)
                 colors = cols.astype(np.uint8)
 
-    # 2) Separate "points" and "colors"/"rgb"
     if points is None:
         if "points" in out:
             pts = out["points"]
@@ -93,11 +79,8 @@ def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
             if pts.ndim == 3 and pts.shape[0] == 1:
                 pts = pts[0]
             pts_np = pts.detach().cpu().float().numpy()
-            if pts_np.shape[-1] != 3:
-                raise RuntimeError(f"Expected points last dim=3, got {pts_np.shape}")
             points = pts_np
 
-            # colors
             col_tensor = None
             for k in ["colors", "rgb", "point_colors"]:
                 if k in out:
@@ -112,40 +95,26 @@ def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
                     col_np = (col_np * 255.0).clip(0, 255)
                 colors = col_np.astype(np.uint8)
 
-    # 3) If still no colors, default to white
-    if points is not None and colors is None:
-        colors = np.full_like(points, 255, dtype=np.uint8)
-
     if points is None:
-        raise RuntimeError(
-            f"Could not find point cloud in MoGe output; keys: {list(out.keys())}"
-        )
+        raise RuntimeError(f"Could not find point cloud in MoGe output")
 
-    # ensure 2D
     points = points.reshape(-1, 3)
-    colors = colors.reshape(-1, 3)
+    if colors is None:
+        colors = np.full_like(points, 255, dtype=np.uint8)
+    else:
+        colors = colors.reshape(-1, 3)
 
     n = points.shape[0]
     print("MoGe point count:", n)
-
-    # sanity check: bail if the model gave us basically nothing
     if n < 100:
-        raise RuntimeError(f"MoGe returned too few points (N={n}), refusing to write bogus PLY.")
+        raise RuntimeError(f"Too few points (N={n}), refusing to export")
 
     return points, colors
 
-# ---------- Helper: write PLY into memory ----------
 
 def pointcloud_to_ply_bytes(points: np.ndarray, colors: np.ndarray) -> bytes:
-    """
-    Create an ASCII PLY file as bytes from points & colors.
-
-    points: (N,3) float32
-    colors: (N,3) uint8
-    """
     n = points.shape[0]
-    pts = points.reshape(-1, 3)
-    cols = colors.reshape(-1, 3)
+    print("Writing PLY with", n, "points")
 
     header = f"""ply
 format ascii 1.0
@@ -160,53 +129,118 @@ end_header
 """
     lines = []
     for i in range(n):
-        x, y, z = pts[i]
-        r, g, b = cols[i]
+        x, y, z = points[i]
+        r, g, b = colors[i]
         lines.append(f"{x:.6f} {y:.6f} {z:.6f} {int(r)} {int(g)} {int(b)}")
 
     body = "\n".join(lines) + "\n"
-    ply_str = header + body
-    return ply_str.encode("utf-8")
+    return (header + body).encode("utf-8")
 
 
-# ---------- Gradio inference function ----------
-
-def infer_and_export_ply(image: np.ndarray) -> gr.File:
+def pointcloud_to_mesh_glb_bytes(points: np.ndarray, colors: np.ndarray) -> bytes:
     """
-    Gradio callable: takes an RGB image, returns a PLY file.
+    Build a surface mesh from the point cloud using Poisson reconstruction,
+    transfer colors from points to mesh vertices via nearest neighbor, and
+    export as GLB with vertex colors.
     """
+    print("Building mesh from point cloud for GLB export")
+    # Optional: downsample for speed
+    max_points = 50000
+    if points.shape[0] > max_points:
+        idx = np.random.choice(points.shape[0], max_points, replace=False)
+        pts_ds = points[idx]
+        cols_ds = colors[idx]
+    else:
+        pts_ds = points
+        cols_ds = colors
+
+    # Open3D point cloud
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(pts_ds.astype(np.float64))
+    pcd.colors = o3d.utility.Vector3dVector((cols_ds / 255.0).astype(np.float64))
+
+    # Poisson reconstruction
+    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
+        pcd, depth=8
+    )
+
+    # Remove low-density vertices (optional cleanup)
+    densities = np.asarray(densities)
+    density_thresh = np.quantile(densities, 0.05)
+    vertices_to_keep = densities > density_thresh
+    mesh = mesh.select_by_index(np.where(vertices_to_keep)[0])
+
+    mesh.remove_duplicated_vertices()
+    mesh.remove_degenerate_triangles()
+    mesh.remove_duplicated_triangles()
+    mesh.remove_non_manifold_edges()
+
+    verts = np.asarray(mesh.vertices)
+    faces = np.asarray(mesh.triangles)
+    print("Mesh verts:", verts.shape, "faces:", faces.shape)
+
+    if verts.shape[0] == 0 or faces.shape[0] == 0:
+        raise RuntimeError("Mesh reconstruction failed; got empty mesh")
+
+    # Transfer colors from original (downsampled) cloud to mesh vertices
+    pcd_tree = o3d.geometry.KDTreeFlann(pcd)
+    vert_colors = []
+    for v in verts:
+        _, idx, _ = pcd_tree.search_knn_vector_3d(v, 1)
+        vert_colors.append(np.asarray(pcd.colors)[idx[0]])
+    vert_colors = np.stack(vert_colors, axis=0)  # (V,3) in [0,1]
+
+    # Convert to trimesh for GLB export
+    tm = trimesh.Trimesh(
+        vertices=verts,
+        faces=faces,
+        vertex_colors=(vert_colors * 255.0).astype(np.uint8),
+        process=False,
+    )
+
+    glb_bytes = tm.export(file_type="glb")
+    if isinstance(glb_bytes, str):
+        glb_bytes = glb_bytes.encode("utf-8")
+    return glb_bytes
+
+
+def infer_and_export_files(image: np.ndarray):
     if image is None:
         raise gr.Error("Please upload an image.")
 
-    # image arrives in HxWx3 RGB (uint8) from gr.Image
     points, colors = run_moge_on_image(image)
 
+    # PLY
     ply_bytes = pointcloud_to_ply_bytes(points, colors)
-
-    # Save to a temporary file path Gradio can return
-    tmp_path = "output.ply"
-    with open(tmp_path, "wb") as f:
+    ply_path = "output.ply"
+    with open(ply_path, "wb") as f:
         f.write(ply_bytes)
 
-    return tmp_path
+    # GLB
+    glb_bytes = pointcloud_to_mesh_glb_bytes(points, colors)
+    glb_path = "output.glb"
+    with open(glb_path, "wb") as f:
+        f.write(glb_bytes)
 
+    return ply_path, glb_path
 
-# ---------- Gradio app ----------
 
-title = "MoGe 3D Reconstruction → Textured Point Cloud (PLY)"
+title = "MoGe 3D Reconstruction → PLY + GLB"
 description = (
-    "Upload an image and get a colored point cloud (PLY) reconstructed by MoGe. "
-    "You can download the PLY and render it in Three.js or any 3D viewer."
+    "Upload an image. MoGe reconstructs a 3D point cloud, which is exported as PLY "
+    "and meshed into a colored GLB suitable for Three.js."
 )
 
 demo = gr.Interface(
-    fn=infer_and_export_ply,
+    fn=infer_and_export_files,
     inputs=gr.Image(type="numpy", label="Input image"),
-    outputs=gr.File(label="Download PLY"),
+    outputs=[
+        gr.File(label="Download PLY (point cloud)"),
+        gr.File(label="Download GLB (colored mesh)"),
+    ],
     title=title,
     description=description,
 )
 
-# Expose for HF Spaces
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)

requirements.txt

@@ -3,4 +3,6 @@ torchvision
 numpy
 opencv-python
 Pillow
+trimesh
+open3d
 git+https://github.com/microsoft/MoGe.git