Update app.py

app.py

@@ -7,89 +7,94 @@ from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 import tempfile
 import os
 
-# --- 1. DA3 MODEL SETUP ---
+# --- DA3 SETTINGS ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-# Depth Anything V3 Small - Higher precision for 2026 workflows
-CHECKPOINT = "depth-anything/DA3-Small" 
+CHECKPOINT = "depth-anything/da3-small" 
 
 processor = AutoImageProcessor.from_pretrained(CHECKPOINT)
 model = AutoModelForDepthEstimation.from_pretrained(CHECKPOINT).to(DEVICE)
 
-def process_da3_to_mesh(input_image):
+def process_textured_mesh(input_image):
     if input_image is None:
         return None, None
     
-    # Resize for processing speed
-    input_image.thumbnail((1024, 1024))
-    
-    # --- 2. V3 DEPTH INFERENCE ---
+    # 1. GENERATE DEPTH
     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
-    rgb_colors = np.array(input_image).reshape(-1, 3) / 255.0 
+    step = 2
+    x, y = np.meshgrid(np.arange(0, width, step), np.arange(0, height, step))
     
-    # --- 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 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
     
-    # 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. ADVANCED MESHING (POISSON RECONSTRUCTION) ---
-    pcd = o3d.geometry.PointCloud()
-    pcd.points = o3d.utility.Vector3dVector(points)
-    pcd.colors = o3d.utility.Vector3dVector(rgb_colors)
+    rows, cols, _ = points.shape
     
-    # 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.]))
+    # 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)
 
-    # 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)
+    # 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])
 
-    # --- 5. FINALIZE & EXPORT ---
-    mesh.translate(-mesh.get_center()) # FORCE CENTER
+    # 4. CONSTRUCT MESH
+    mesh = o3d.geometry.TriangleMesh()
+    mesh.vertices = o3d.utility.Vector3dVector(vertices)
+    mesh.triangles = o3d.utility.Vector3iVector(np.array(faces))
     
+    # Assign UVs (Open3D expects UVs per triangle vertex, so we tile them)
+    mesh.triangle_uvs = o3d.utility.Vector2dVector(np.tile(uvs, (3, 1)))
+
+    # 5. EXPORT
     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)
+    mesh_path = os.path.join(temp_dir, "model.obj")
+    texture_path = os.path.join(temp_dir, "texture.png")
+    
+    # Save image as texture
+    input_image.save(texture_path)
+    
+    # Save OBJ
+    o3d.io.write_triangle_mesh(mesh_path, mesh)
     
-    return output_path, output_path
+    # 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
 
-# --- 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.")
+# --- UI ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 🎭 DA3 Textured 3D Mesh")
     
     with gr.Row():
         with gr.Column():
-            img_in = gr.Image(type="pil", label="Source Image")
-            btn = gr.Button("🔨 Generate High-Res Mesh", variant="primary")
+            img_in = gr.Image(type="pil", label="Input")
+            btn = gr.Button("🔨 Generate Mesh", variant="primary")
+        
         with gr.Column():
-            v3d = gr.Model3D(
-                label="3D Mesh Preview", 
-                display_mode="solid", 
-                camera_position=(0, 90, 15)
-            )
-            dl = gr.DownloadButton("💾 Download Colored PLY")
+            # 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")
 
-    btn.click(fn=process_da3_to_mesh, inputs=[img_in], outputs=[v3d, dl])
+    btn.click(fn=process_textured_mesh, inputs=[img_in], outputs=[v3d, dl])
 
 demo.launch()