Update app.py

app.py

@@ -7,21 +7,26 @@ from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 import tempfile
 import os
 
-# --- 1. MODEL SETUP ---
+# --- 1. DA3 MODEL SETUP ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-CHECKPOINT = "depth-anything/Depth-Anything-V2-Small-hf"
+# 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_to_3d(input_image):
+def process_da3_to_mesh(input_image):
     if input_image is None:
         return None, None
     
-    # --- 2. DEPTH & COLOR ---
+    # 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],
@@ -29,65 +34,62 @@ def process_to_3d(input_image):
         ).squeeze().cpu().numpy()
 
     width, height = input_image.size
-    # Explicitly pull RGB and normalize for Open3D
     rgb_colors = np.array(input_image).reshape(-1, 3) / 255.0 
     
-    # --- 3. THE "PERFECT FOCUS" PROJECTION ---
+    # --- 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
     
-    # Normalize depth to a 0-5 range
-    z = (depth.flatten() / (depth.max() + 1e-5)) * 5.0
-    
-    # Normalize X and Y to a -5 to 5 range (Centering)
-    aspect_ratio = width / height
-    x_centered = ((x.flatten() / width) - 0.5) * 10.0 * aspect_ratio
+    # 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)
 
-    # --- 4. OPEN3D COLOR BINDING ---
+    # --- 4. ADVANCED MESHING (POISSON RECONSTRUCTION) ---
     pcd = o3d.geometry.PointCloud()
     pcd.points = o3d.utility.Vector3dVector(points)
     pcd.colors = o3d.utility.Vector3dVector(rgb_colors)
+    
+    # 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.]))
 
-    # Statistical Outlier Removal (Cleans up "ghost" points that mess up camera focus)
-    pcd, _ = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=2.0)
-
-    # Center the model's pivot point at exactly 0,0,0
-    pcd.translate(-pcd.get_center())
-
-    # Voxelize to make points "solid"
-    pcd = pcd.voxel_down_sample(voxel_size=0.03)
+    # 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)
+    
+    # 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. SAVE BINARY PLY ---
+    # --- 5. FINALIZE & EXPORT ---
+    mesh.translate(-mesh.get_center()) # FORCE CENTER
+    
     temp_dir = tempfile.gettempdir()
-    output_path = os.path.join(temp_dir, "model.ply")
-    # write_ascii=False ensures color data is encoded in a way Gradio understands
-    o3d.io.write_point_cloud(output_path, pcd, write_ascii=False)
+    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)
     
     return output_path, output_path
 
-# --- 6. UI WITH CAMERA FOCUS ---
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧊 High-Fidelity 3D Splat View")
-    gr.Markdown("The camera is pre-focused on the object. Click and drag to rotate.")
+# --- 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")
-            btn = gr.Button("🔨 Generate & Focus", variant="primary")
-        
+            img_in = gr.Image(type="pil", label="Source Image")
+            btn = gr.Button("🔨 Generate High-Res Mesh", variant="primary")
         with gr.Column():
-            # camera_position: (azimuth, elevation, distance)
-            # Distance 12 is perfect for our -5 to 5 coordinate range
             v3d = gr.Model3D(
-                label="3D Viewport",
+                label="3D Mesh Preview", 
                 display_mode="solid", 
-                camera_position=(0, 90, 12), 
-                clear_color=(0.1, 0.1, 0.1, 1.0)
+                camera_position=(0, 90, 15)
             )
-            dl = gr.DownloadButton("💾 Download Colored .PLY")
+            dl = gr.DownloadButton("💾 Download Colored PLY")
 
-    btn.click(fn=process_to_3d, inputs=[img_in], outputs=[v3d, dl])
+    btn.click(fn=process_da3_to_mesh, inputs=[img_in], outputs=[v3d, dl])
 
 demo.launch()