Upload app.py

app.py

@@ -39,65 +39,100 @@ def run_moge_on_image(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
     image: HxWx3 RGB uint8 numpy array.
 
     Returns:
-        points: (N, 3) float32 XYZ in some model-defined coordinates
-        colors: (N, 3) uint8 RGB in [0, 255]
+        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)
 
-    # NOTE: This assumes MoGeModel.infer returns something like:
-    #   {"points": (1, N, 3), "colors": (1, N, 3)} or similar.
-    # You may need to adapt this part to the actual MoGe API.
+    # --- Run MoGe ---
     out = MODEL.infer(tensor)
 
-    # ----- Adapt this based on the actual return structure -----
-    #
-    # Common patterns:
-    #   out["points"]: (B, N, 3) point coordinates
-    #   out["colors"]: (B, N, 3) colors in [0, 1] or [0, 255]
-    #
-    # If your actual keys or shapes are different, adjust here.
-
-    if "points" in out:
-        points = out["points"]
-    elif "point_cloud" in out:
-        points = out["point_cloud"]
-    else:
-        raise RuntimeError(f"Cannot find point cloud in MoGe output keys: {list(out.keys())}")
-
-    # remove batch dim
-    if points.ndim == 3:
-        points = points[0]
-
-    points = points.detach().cpu().float().numpy()  # (N,3)
-
-    # Try to get colors if available, else default to white
-    colors_raw = None
-    for k in ["colors", "rgb", "point_colors"]:
-        if k in out:
-            colors_raw = out[k]
-            break
-
-    if colors_raw is not None:
-        if colors_raw.ndim == 3:
-            colors_raw = colors_raw[0]
-        colors_np = colors_raw.detach().cpu().float().numpy()
-        # Normalize to [0,255] if necessary
-        if colors_np.max() <= 1.0:
-            colors_np = (colors_np * 255.0).clip(0, 255)
-        colors = colors_np.astype(np.uint8)
-    else:
-        # fallback: all white
+    # --- 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 ---
+
+    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)
+            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)
+
+    # 2) Separate "points" and "colors"/"rgb"
+    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()
+            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:
+                    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)
+
+    # 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)
 
-    # Ensure shapes
-    assert points.shape[-1] == 3, f"Expected points (N,3), got {points.shape}"
-    assert colors.shape[-1] == 3, f"Expected colors (N,3), got {colors.shape}"
+    if points is None:
+        raise RuntimeError(
+            f"Could not find point cloud in MoGe output; keys: {list(out.keys())}"
+        )
 
-    return points, colors
+    # ensure 2D
+    points = points.reshape(-1, 3)
+    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.")
+
+    return points, colors
 
 # ---------- Helper: write PLY into memory ----------