Upload 2 files

app.py

@@ -1,4 +1,3 @@
-import io
 import os
 from typing import Tuple
 
@@ -6,17 +5,26 @@ import gradio as gr
 import numpy as np
 import torch
 import cv2
-import open3d as o3d
-import trimesh
 
 from moge.model.v2 import MoGeModel
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 
+# ---------- Model setup ----------
+
 @torch.no_grad()
 def load_model() -> MoGeModel:
+    """
+    Load the mesh-capable MoGe model.
+
+    NOTE:
+    - If there is a dedicated mesh checkpoint (e.g. "Ruicheng/moge-2-vitl-mesh"),
+      use that ID here.
+    - If not, keep the normal one and use the mesh reconstruction API on it.
+    """
     print(f"Loading MoGe model on device: {DEVICE}")
+    # If there is a mesh-specific checkpoint, change this string accordingly.
     model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal")
     model = model.to(DEVICE)
     model.eval()
@@ -26,15 +34,18 @@ def load_model() -> MoGeModel:
 MODEL = load_model()
 
 
+# ---------- Helper: run MoGe mesh reconstruction ----------
+
 @torch.no_grad()
-def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+def run_moge_mesh(image: np.ndarray) -> bytes:
     """
     image: HxWx3 RGB uint8 numpy array.
 
     Returns:
-        points: (N, 3) float32 XYZ
-        colors: (N, 3) uint8 RGB
+        glb_bytes: binary GLB data with texture baked, resolution ~256.
     """
+
+    # Convert to float [0,1], CHW, batch
     img = image.astype(np.float32) / 255.0
     tensor = (
         torch.from_numpy(img)
@@ -43,226 +54,102 @@ def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
         .to(DEVICE)  # (1,3,H,W)
     )
 
-    out = MODEL.infer(tensor)
-
-    print("MoGe output keys:", list(out.keys()))
-
-    # 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
-
-    if "pcd" in out:
-        pcd = out["pcd"]
-        if pcd.ndim == 3 and pcd.shape[-1] >= 3:
-            if pcd.shape[0] == 1:
-                pcd = pcd[0]
-            pcd_np = pcd.detach().cpu().float().numpy()
-            points = pcd_np[:, :3]
-            if pcd_np.shape[1] >= 6:
-                cols = pcd_np[:, 3:6]
-                if cols.max() <= 1.0:
-                    cols = (cols * 255.0).clip(0, 255)
-                colors = cols.astype(np.uint8)
-
-    if points is None:
-        if "points" in out:
-            pts = out["points"]
-        elif "point_cloud" in out:
-            pts = out["point_cloud"]
-        else:
-            pts = None
-
-        if pts is not None:
-            if pts.ndim == 3 and pts.shape[0] == 1:
-                pts = pts[0]
-            pts_np = pts.detach().cpu().float().numpy()
-            points = pts_np
-
-            col_tensor = None
-            for k in ["colors", "rgb", "point_colors"]:
-                if k in out:
-                    col_tensor = out[k]
-                    break
-
-            if col_tensor is not None:
-                if col_tensor.ndim == 3 and col_tensor.shape[0] == 1:
-                    col_tensor = col_tensor[0]
-                col_np = col_tensor.detach().cpu().float().numpy()
-                if col_np.max() <= 1.0:
-                    col_np = (col_np * 255.0).clip(0, 255)
-                colors = col_np.astype(np.uint8)
-
-    if points is None:
-        raise RuntimeError(f"Could not find point cloud in MoGe output")
-
-    points = points.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)
-    if n < 100:
-        raise RuntimeError(f"Too few points (N={n}), refusing to export")
-
-    return points, colors
-
-
-def pointcloud_to_ply_bytes(points: np.ndarray, colors: np.ndarray) -> bytes:
-    n = points.shape[0]
-    print("Writing PLY with", n, "points")
-
-    header = f"""ply
-format ascii 1.0
-element vertex {n}
-property float x
-property float y
-property float z
-property uchar red
-property uchar green
-property uchar blue
-end_header
-"""
-    lines = []
-    for i in range(n):
-        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"
-    return (header + body).encode("utf-8")
-
-
-def pointcloud_to_mesh_glb_bytes(points: np.ndarray, colors: np.ndarray) -> bytes:
-    """
-    Build a surface mesh from the point cloud using Poisson reconstruction,
-    denoise it, 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")
-
-    # Basic normalization: center the cloud to reduce numeric issues
-    center = points.mean(axis=0, keepdims=True)
-    pts_norm = points - center
-
-    # Optional: downsample for speed
-    pcd = o3d.geometry.PointCloud()
-    pcd.points = o3d.utility.Vector3dVector(pts_norm.astype(np.float64))
-    pcd.colors = o3d.utility.Vector3dVector((colors / 255.0).astype(np.float64))
-
-    # Voxel downsample: tweak voxel_size depending on MoGe scale
-    voxel_size = float(np.linalg.norm(pts_norm.max(axis=0) - pts_norm.min(axis=0)) / 128.0)
-    print("Voxel size:", voxel_size)
-    if voxel_size > 0:
-        pcd = pcd.voxel_down_sample(voxel_size=voxel_size)
-
-    print("After downsample:", np.asarray(pcd.points).shape[0], "points")
-
-    # Remove obvious outliers (radius-based)
-    print("Removing outliers...")
-    try:
-        pcd, _ = pcd.remove_radius_outlier(nb_points=20, radius=voxel_size * 3.0)
-    except Exception as e:
-        print("Outlier removal failed:", e)
-
-    print("Estimating normals...")
-    pcd.estimate_normals(
-        search_param=o3d.geometry.KDTreeSearchParamKNN(knn=30)
+    # ---- IMPORTANT PART: call the mesh reconstruction API ----
+    #
+    # You need to adjust THIS CALL to match the actual MoGe code.
+    # Look for something like:
+    #   - MODEL.reconstruct_mesh(...)
+    #   - MODEL.mesh_reconstruct(...)
+    #   - MODEL.infer_mesh(...)
+    #
+    # And for arguments, look for:
+    #   - mesh_resolution / grid_resolution
+    #   - texture_size / tex_size
+    #   - enable_texture / with_texture
+    #
+    # Below is a TEMPLATE that you should modify once you've checked the repo.
+
+    # TEMPLATE call – this will almost certainly need renaming:
+    result = MODEL.reconstruct_mesh(
+        tensor,
+        mesh_resolution=256,     # 256^3 grid or equivalent
+        texture_size=256,        # 256x256 texture
+        enable_texture=True,     # or with_texture=True, etc.
     )
 
-    # Orient normals consistently
-    try:
-        pcd.orient_normals_consistent_tangent_plane(30)
-    except Exception as e:
-        print("Normal orientation failed:", e)
-
-    # Poisson reconstruction
-    print("Running Poisson reconstruction...")
-    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
-        pcd, depth=8
+    # ---- Inspect result structure (one-time debugging) ----
+    # While debugging, you can keep these prints to see keys in Space logs:
+    print("MoGe mesh result keys:", list(result.keys()))
+
+    # Common patterns:
+    #  1) result["glb"] -> raw GLB bytes
+    #  2) result["mesh"] -> mesh object (trimesh / internal) with export method
+
+    # Case 1: GLB bytes directly
+    if "glb" in result:
+        glb_bytes = result["glb"]
+        if isinstance(glb_bytes, str):
+            glb_bytes = glb_bytes.encode("utf-8")
+        return glb_bytes
+
+    # Case 2: mesh object with export method
+    if "mesh" in result:
+        mesh = result["mesh"]
+        # If MoGe mesh exposes something like `to_glb(texture=..., texture_size=256)`:
+        if hasattr(mesh, "to_glb"):
+            tex = result.get("texture", None)
+            if tex is not None:
+                glb_bytes = mesh.to_glb(texture=tex, texture_size=256)
+            else:
+                glb_bytes = mesh.to_glb(texture_size=256)
+            if isinstance(glb_bytes, str):
+                glb_bytes = glb_bytes.encode("utf-8")
+            return glb_bytes
+
+        # Or if it expects file export:
+        if hasattr(mesh, "export"):
+            tmp_path = "output.glb"
+            tex = result.get("texture", None)
+            if tex is not None:
+                # This is pseudocode – adapt to the actual mesh.export signature.
+                mesh.export(tmp_path, texture=tex, texture_size=256)
+            else:
+                mesh.export(tmp_path)
+            with open(tmp_path, "rb") as f:
+                return f.read()
+
+    raise RuntimeError(
+        f"Unsupported MoGe mesh result structure: keys={list(result.keys())}"
     )
 
-    densities = np.asarray(densities)
-    # Keep higher-density areas: cuts away wispy boundary fog
-    density_thresh = np.quantile(densities, 0.1)
-    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 filtered point cloud -> mesh vertices
-    print("Transferring vertex colors...")
-    pcd_tree = o3d.geometry.KDTreeFlann(pcd)
-    pcd_colors_np = np.asarray(pcd.colors)
-    vert_colors = []
-    for v in verts:
-        _, idx, _ = pcd_tree.search_knn_vector_3d(v, 1)
-        vert_colors.append(pcd_colors_np[idx[0]])
-    vert_colors = np.stack(vert_colors, axis=0)  # (V,3) in [0,1]
-
-    # Undo centering so the mesh is in original coordinates
-    verts = verts + center
-
-    # 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):
+# ---------- Gradio inference function ----------
+
+def infer_and_export_glb(image: np.ndarray):
     if image is None:
         raise gr.Error("Please upload an image.")
 
-    points, colors = run_moge_on_image(image)
-
-    # PLY
-    ply_bytes = pointcloud_to_ply_bytes(points, colors)
-    ply_path = "output.ply"
-    with open(ply_path, "wb") as f:
-        f.write(ply_bytes)
+    glb_bytes = run_moge_mesh(image)
 
-    # 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
+    return glb_path
+
 
+# ---------- Gradio app ----------
 
-title = "MoGe 3D Reconstruction → PLY + GLB"
+title = "MoGe 3D Reconstruction → Textured GLB (256)"
 description = (
-    "Upload an image. MoGe reconstructs a 3D point cloud, which is exported as PLY "
-    "and meshed into a colored GLB suitable for Three.js."
+    "Upload an image. MoGe reconstructs a textured 3D mesh and exports it as a GLB "
+    "with a ~256x256 texture."
 )
 
 demo = gr.Interface(
-    fn=infer_and_export_files,
+    fn=infer_and_export_glb,
     inputs=gr.Image(type="numpy", label="Input image"),
-    outputs=[
-        gr.File(label="Download PLY (point cloud)"),
-        gr.File(label="Download GLB (colored mesh)"),
-    ],
+    outputs=gr.File(label="Download GLB (textured mesh)"),
     title=title,
     description=description,
 )

requirements.txt

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