Update app.py

app.py

@@ -1,5 +1,4 @@
-# Minimal 2D -> 3D (GLPN + Open3D) — fixed geometry + reliable viewer
-
+# 2D -> 3D (GLPN + Open3D) — Mesh-only (Poisson) output
 import os, time, traceback
 from pathlib import Path
 import numpy as np
@@ -9,11 +8,11 @@ import gradio as gr
 import open3d as o3d
 from transformers import GLPNForDepthEstimation, GLPNImageProcessor
 
-# Quiet threads warnings in Spaces
+# Quiet HF threads warnings
 os.environ.setdefault("OMP_NUM_THREADS", "1")
 os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
 
-# -------- Model / device --------
+# ---- Model / device ----
 DEVICE = torch.device(
     "cuda" if torch.cuda.is_available()
     else ("mps" if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available() else "cpu")
@@ -22,7 +21,7 @@ MODEL_ID = "vinvino02/glpn-nyu"
 PROCESSOR = GLPNImageProcessor.from_pretrained(MODEL_ID)
 MODEL = GLPNForDepthEstimation.from_pretrained(MODEL_ID).to(DEVICE).eval()
 
-# -------- Helpers (simple; no knobs) --------
+# ---- Helpers ----
 def _resize_main(pil_img: Image.Image):
     new_h = max(32, min(pil_img.height, 480))
     new_h -= new_h % 32
@@ -31,32 +30,26 @@ def _resize_main(pil_img: Image.Image):
 
 @torch.inference_mode()
 def _depth_pred_float(pil_img: Image.Image) -> np.ndarray:
-    """Predict depth, upsample to original size (float32, arbitrary scale)."""
     resized, (W, H) = _resize_main(pil_img)
     inputs = PROCESSOR(images=resized, return_tensors="pt")
     inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
     out = MODEL(**inputs).predicted_depth  # [1, h, w]
     up = torch.nn.functional.interpolate(out.unsqueeze(1), size=(H, W), mode="bicubic", align_corners=False).squeeze(1)
-    d = up[0].detach().float().cpu().numpy()
-    return d
+    return up[0].detach().float().cpu().numpy()
 
 def _depth_preview_u8(d: np.ndarray) -> Image.Image:
     d = d - d.min()
     mx = float(d.max()) if d.size else 1.0
     if mx <= 0: mx = 1.0
-    u8 = (255.0 * d / mx).astype(np.uint8)
-    return Image.fromarray(u8)
+    return Image.fromarray((255.0 * d / mx).astype(np.uint8))
 
 def _depth_to_metric_meters(d: np.ndarray, near=0.3, far=5.0) -> np.ndarray:
-    """Map arbitrary depth to meters using robust normalization (2–98% percentile)."""
     lo, hi = np.percentile(d, [2.0, 98.0])
     d01 = np.clip((d - lo) / max(hi - lo, 1e-6), 0, 1).astype(np.float32)
-    # 0..1 -> [near, far] meters
     return (near + d01 * (far - near)).astype(np.float32)
 
 def _rgbd_for_open3d(rgb: Image.Image, depth_m: np.ndarray, far=5.0) -> o3d.geometry.RGBDImage:
-    # Open3D expects depth in the same units used by depth_scale denominator.
-    depth_scale = 1000.0  # meters * 1000 (mm-like)
+    depth_scale = 1000.0  # meters * 1000
     depth_o3d = o3d.geometry.Image((depth_m * depth_scale).astype(np.float32))
     color_o3d = o3d.geometry.Image(np.array(rgb.convert("RGB")))
     return o3d.geometry.RGBDImage.create_from_color_and_depth(
@@ -67,11 +60,9 @@ def _rgbd_for_open3d(rgb: Image.Image, depth_m: np.ndarray, far=5.0) -> o3d.geom
 def _pcd_from_rgbd(rgbd: o3d.geometry.RGBDImage) -> o3d.geometry.PointCloud:
     h = np.asarray(rgbd.depth).shape[0]
     w = np.asarray(rgbd.depth).shape[1]
-    # keep your original intrinsics spirit; fx=fy=500, cx=w/2, cy=h/2
     intr = o3d.camera.PinholeCameraIntrinsic(w, h, 500.0, 500.0, w/2.0, h/2.0)
     pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd, intr)
-    # Make orientation sensible for web viewers
-    pcd.transform([[1,0,0,0], [0,-1,0,0], [0,0,-1,0], [0,0,0,1]])
+    pcd.transform([[1,0,0,0],[0,-1,0,0],[0,0,-1,0],[0,0,0,1]])  # upright for web
     return pcd
 
 def _clean_pcd(pcd: o3d.geometry.PointCloud) -> o3d.geometry.PointCloud:
@@ -82,20 +73,10 @@ def _clean_pcd(pcd: o3d.geometry.PointCloud) -> o3d.geometry.PointCloud:
     pcd.orient_normals_consistent_tangent_plane(10)
     return pcd
 
-def _mesh_from_pcd(pcd: o3d.geometry.PointCloud) -> o3d.geometry.TriangleMesh:
+def _poisson_mesh_from_pcd(pcd: o3d.geometry.PointCloud) -> o3d.geometry.TriangleMesh:
     if len(pcd.points) == 0: return o3d.geometry.TriangleMesh()
-    # Ball-Pivoting first (robust), fallback to Poisson
-    try:
-        dists = np.asarray(pcd.compute_nearest_neighbor_distance())
-        avg = float(np.mean(dists)) if dists.size else 0.01
-        radii = [avg*2.0, avg*3.0, avg*4.0]
-        mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(
-            pcd, o3d.utility.DoubleVector(radii)
-        )
-    except Exception:
-        mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=10, n_threads=1)
-    # Rotate like your main.py
-    R = mesh.get_rotation_matrix_from_xyz((np.pi, 0.0, 0.0))
+    mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=10, n_threads=1)
+    R = mesh.get_rotation_matrix_from_xyz((np.pi, 0.0, 0.0))  # match main.py
     mesh.rotate(R, center=(0,0,0))
     mesh.remove_degenerate_triangles(); mesh.remove_duplicated_vertices()
     mesh.remove_non_manifold_edges(); mesh.remove_unreferenced_vertices()
@@ -103,18 +84,16 @@ def _mesh_from_pcd(pcd: o3d.geometry.PointCloud) -> o3d.geometry.TriangleMesh:
     return mesh
 
 def _normalize_for_view(mesh: o3d.geometry.TriangleMesh) -> o3d.geometry.TriangleMesh:
-    """Center + unit scale so it always shows up in the viewer."""
     if len(mesh.vertices) == 0: return mesh
     aabb = mesh.get_axis_aligned_bounding_box()
     c = aabb.get_center()
     mesh = mesh.translate(-c, relative=True)
-    extent = aabb.get_extent()
-    s = 1.0 / max(extent.max(), 1e-6)
+    s = 1.0 / max(aabb.get_extent().max(), 1e-6)
     mesh = mesh.scale(s, center=(0,0,0))
     mesh.compute_vertex_normals()
     return mesh
 
-# -------- Pipeline (no options) --------
+# ---- Pipeline (mesh-only output) ----
 def run(image: Image.Image):
     logs = []
     t0 = time.time()
@@ -125,42 +104,39 @@ def run(image: Image.Image):
         image = image.convert("RGB")
         logs.append("1) Predicting depth (GLPN)…")
         d_pred = _depth_pred_float(image)
-        depth_preview = _depth_preview_u8(d_pred)   # for UI only
+        depth_preview = _depth_preview_u8(d_pred)
 
         logs.append("2) Mapping to meters (0.3–5.0 m)…")
         d_m = _depth_to_metric_meters(d_pred, near=0.3, far=5.0)
 
-        logs.append("3) Building RGBD and point cloud…")
+        logs.append("3) RGBD -> PCD -> Poisson mesh…")
         rgbd = _rgbd_for_open3d(image, d_m, far=5.0)
-        pcd = _pcd_from_rgbd(rgbd)
-        pcd = _clean_pcd(pcd)
-
-        logs.append("4) Meshing + rotate…")
-        mesh = _mesh_from_pcd(pcd)
+        pcd  = _pcd_from_rgbd(rgbd)
+        pcd  = _clean_pcd(pcd)
+        mesh = _poisson_mesh_from_pcd(pcd)
 
-        # Save raw assets (metric)
+        # Save ONLY the mesh
         out = Path("outputs"); out.mkdir(parents=True, exist_ok=True)
-        pcd_path   = str(out / "point_cloud.ply")
         mesh_path  = str(out / "mesh.ply")
-        o3d.io.write_point_cloud(pcd_path, pcd)
         o3d.io.write_triangle_mesh(mesh_path, mesh)
+        logs.append(f"Saved mesh → {mesh_path}")
+        logs.append(f"Mesh stats: Vertices={len(mesh.vertices):,}  Triangles={len(mesh.triangles):,}")
 
-        # Save a viewer copy (normalized so it always displays)
+        # Viewer copy (normalized so it always shows)
         viewer_path = str(out / "mesh_viewer.ply")
         o3d.io.write_triangle_mesh(viewer_path, _normalize_for_view(mesh))
-
-        logs.append("5) Saved outputs.")
         logs.append(f"Done in {time.time()-t0:.1f}s.")
-        return depth_preview, viewer_path, pcd_path, mesh_path, "\n".join(logs)
+
+        return depth_preview, viewer_path, mesh_path, "\n".join(logs)
 
     except Exception as e:
         tb = traceback.format_exc()
         logs.append(f"[ERROR] {e}\n{tb}")
-        return None, None, None, None, "\n".join(logs)
+        return None, None, None, "\n".join(logs)
 
-# -------- UI --------
-with gr.Blocks(title="2D → 3D (GLPN + Open3D) — Minimal & Fixed") as demo:
-    gr.Markdown("### 2D → 3D (GLPN + Open3D)\nUpload → Depth preview → 3D mesh (always visible) → Downloads.")
+# ---- UI ----
+with gr.Blocks(title="2D → 3D (GLPN + Open3D) — Mesh Only") as demo:
+    gr.Markdown("### 2D → 3D — Mesh Only (Poisson)\nUpload → Depth preview → **Triangle mesh** (viewer + PLY download).")
 
     with gr.Row():
         with gr.Column():
@@ -170,12 +146,10 @@ with gr.Blocks(title="2D → 3D (GLPN + Open3D) — Minimal & Fixed") as demo:
 
         with gr.Column():
             depth_img = gr.Image(label="Depth (preview)")
-            model3d   = gr.Model3D(label="3D Mesh (normalized for viewing)", height=520)
-            with gr.Row():
-                pcd_file  = gr.File(label="point_cloud.ply")
-                mesh_file = gr.File(label="mesh.ply")
+            model3d   = gr.Model3D(label="Triangle Mesh (normalized for viewing)", height=520)
+            mesh_file = gr.File(label="mesh.ply")
 
-    btn.click(run, inputs=[inp], outputs=[depth_img, model3d, pcd_file, mesh_file, logs])
+    btn.click(run, inputs=[inp], outputs=[depth_img, model3d, mesh_file, logs])
 
 demo.queue()
 demo.launch(ssr_mode=False)