Update app.py

app.py

@@ -1,100 +1,110 @@
-import gradio as gr
-import torch
-import numpy as np
-import open3d as o3d
-from PIL import Image
-from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 import tempfile
 import os
 
-# --- DA3 SETTINGS ---
+# --- 1. DA3 MODEL SETUP ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-CHECKPOINT = "depth-anything/da3-small" 
+# Depth Anything V3 Small - Higher precision for 2026 workflows
+CHECKPOINT = "depth-anything/DA3-Small" 
 
 processor = AutoImageProcessor.from_pretrained(CHECKPOINT)
 model = AutoModelForDepthEstimation.from_pretrained(CHECKPOINT).to(DEVICE)
 
-def process_textured_mesh(input_image):
+def process_da3_to_mesh(input_image):
     if input_image is None:
         return None, None
     
-    # 1. GENERATE DEPTH
+    # Resize for processing speed
+    input_image.thumbnail((1024, 1024))
+    
+    # --- 2. V3 DEPTH INFERENCE ---
     inputs = processor(images=input_image, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         outputs = model(**inputs)
+        # DA3 provides much sharper depth maps
         depth = torch.nn.functional.interpolate(
             outputs.predicted_depth.unsqueeze(1),
             size=input_image.size[::-1],
             mode="bicubic",
         ).squeeze().cpu().numpy()
 
-    # 2. CREATE TEXTURED GRID
-    # We use a step of 2 to keep the mesh lightweight for the browser
+
+
     width, height = input_image.size
-    step = 2
-    x, y = np.meshgrid(np.arange(0, width, step), np.arange(0, height, step))
+    rgb_colors = np.array(input_image).reshape(-1, 3) / 255.0 
+
     
-    # Normalize Z (depth) and center X, Y in a unit-10 space
-    z = (depth[::step, ::step] / (depth.max() + 1e-5)) * 3.0
-    x_centered = ((x / width) - 0.5) * 10.0 * (width / height)
-    y_centered = (0.5 - (y / height)) * 10.0
+    # --- 3. NORMALIZED 3D PROJECTION ---
+    x, y = np.meshgrid(np.arange(width), np.arange(height))
+    # DA3 depth is more linear; we scale it for a natural 3D look
+    z = (depth.flatten() / (depth.max() + 1e-5)) * 4.0
     
+    # Center everything in the 'Unit 10' viewing box
+    x_centered = ((x.flatten() / width) - 0.5) * 10.0 * (width / height)
+    y_centered = (0.5 - (y.flatten() / height)) * 10.0 
     points = np.stack((x_centered, y_centered, z), axis=-1)
-    rows, cols, _ = points.shape
-    
-    # 3. VERTICES & UV MAPPING
-    vertices = points.reshape(-1, 3)
-    # UVs map the image (0-1 range) to the vertices
-    uvs = np.stack((x / width, 1.0 - (y / height)), axis=-1).reshape(-1, 2)
-
-    # Build Triangles
-    faces = []
-    for i in range(rows - 1):
-        for j in range(cols - 1):
-            v0 = i * cols + j
-            v1 = v0 + 1
-            v2 = (i + 1) * cols + j
-            v3 = v2 + 1
-            faces.append([v0, v2, v1])
-            faces.append([v1, v2, v3])
-
-    # 4. CONSTRUCT MESH
-    mesh = o3d.geometry.TriangleMesh()
-    mesh.vertices = o3d.utility.Vector3dVector(vertices)
-    mesh.triangles = o3d.utility.Vector3iVector(np.array(faces))
+
+    # --- 4. ADVANCED MESHING (POISSON RECONSTRUCTION) ---
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+    pcd.colors = o3d.utility.Vector3dVector(rgb_colors)
     
-    # Assign UVs (Open3D expects UVs per triangle vertex, so we tile them)
-    mesh.triangle_uvs = o3d.utility.Vector2dVector(np.tile(uvs, (3, 1)))
+    # Estimate Normals - DA3 needs higher search radius for its high-detail output
+    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.2, max_nn=50))
+    pcd.orient_normals_towards_camera_location(camera_location=np.array([0., 0., 15.]))
 
-    # 5. EXPORT
-    temp_dir = tempfile.gettempdir()
-    mesh_path = os.path.join(temp_dir, "model.obj")
-    texture_path = os.path.join(temp_dir, "texture.png")
+
+    # Poisson Surface Reconstruction creates a watertight "solid" shell
+    # depth=8 or 9 is the sweet spot for detail vs speed
+    mesh, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=9)
     
-    # Save image as texture
-    input_image.save(texture_path)
+    # Clean up the mesh (Poisson creates a 'bubble' we need to trim)
+    vertices_to_remove = densities < np.quantile(densities, 0.1)
+    mesh.remove_vertices_by_mask(vertices_to_remove)
+
+
+
+
+    # --- 5. FINALIZE & EXPORT ---
+    mesh.translate(-mesh.get_center()) # FORCE CENTER
+
+
     
-    # Save OBJ
-    o3d.io.write_triangle_mesh(mesh_path, mesh)
+
+
+
+
+    temp_dir = tempfile.gettempdir()
+    output_path = os.path.join(temp_dir, "da3_mesh.ply")
+    # Binary PLY for Blender Color Compatibility
+    o3d.io.write_triangle_mesh(output_path, mesh, write_ascii=False)
+
+
+
+
+
     
-    # To see textures in some viewers, we return the OBJ. 
-    # In Blender, you'll simply load this texture.png onto the model.
-    return mesh_path, mesh_path
+    return output_path, output_path
 
-# --- UI ---
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎭 DA3 Textured 3D Mesh")
+
+
+# --- 6. UI ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🧊 Depth Anything V3: Mesh Engine")
+    gr.Markdown("Using the 2026 DA3 architecture for high-fidelity 3D reconstruction.")
     
     with gr.Row():
         with gr.Column():
-            img_in = gr.Image(type="pil", label="Input")
-            btn = gr.Button("🔨 Generate Mesh", variant="primary")
-        
+            img_in = gr.Image(type="pil", label="Source Image")
+            btn = gr.Button("🔨 Generate High-Res Mesh", variant="primary")
+
         with gr.Column():
-            # Gradio 5.0+ focuses on the center (0,0,0) automatically
-            v3d = gr.Model3D(label="3D Preview", camera_position=(0, 90, 15))
-            dl = gr.DownloadButton("💾 Download OBJ + PNG")
+            v3d = gr.Model3D(
+                label="3D Mesh Preview", 
+                display_mode="solid", 
+                camera_position=(0, 90, 15)
+            )
+            dl = gr.DownloadButton("💾 Download Colored PLY")
 
-    btn.click(fn=process_textured_mesh, inputs=[img_in], outputs=[v3d, dl])
+    btn.click(fn=process_da3_to_mesh, inputs=[img_in], outputs=[v3d, dl])
 
 demo.launch()