Update app.py

app.py

@@ -7,7 +7,7 @@ from transformers import AutoImageProcessor, AutoModelForDepthEstimation
 import tempfile
 import os
 
-# --- 1. SETTINGS ---
+# --- 1. MODEL SETUP ---
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 CHECKPOINT = "depth-anything/Depth-Anything-V2-Small-hf"
 
@@ -18,11 +18,7 @@ def process_to_3d(input_image):
     if input_image is None:
         return None, None
     
-    # Resize image to a manageable size for 3D viewing if too large
-    if max(input_image.size) > 1024:
-        input_image.thumbnail((1024, 1024))
-
-    # --- 2. DEPTH INFERENCE ---
+    # --- 2. DEPTH & COLOR ---
     inputs = processor(images=input_image, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         outputs = model(**inputs)
@@ -32,60 +28,65 @@ def process_to_3d(input_image):
             mode="bicubic",
         ).squeeze().cpu().numpy()
 
-    # --- 3. COLOR & COORDINATE CALCULATION ---
     width, height = input_image.size
-    rgb = np.array(input_image).reshape(-1, 3) / 255.0  # Normalize to 0-1 for O3D
+    # Explicitly pull RGB and normalize for Open3D
+    rgb_colors = np.array(input_image).reshape(-1, 3) / 255.0 
     
-    # Create normalized grid
+    # --- 3. THE "PERFECT FOCUS" PROJECTION ---
     x, y = np.meshgrid(np.arange(width), np.arange(height))
     
-    # Flatten and project to 3D
-    # Scale depth (z) significantly down so it doesn't "stretch" too far back
-    z = (depth.flatten() / depth.max()) * 5.0 
-    x = (x.flatten() - width / 2) / (width / 5.0)
-    y = (height / 2 - y.flatten()) / (height / 5.0) # Invert Y for correct orientation
+    # 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
+    y_centered = (0.5 - (y.flatten() / height)) * 10.0 
 
-    points = np.stack((x, y, z), axis=-1)
+    points = np.stack((x_centered, y_centered, z), axis=-1)
 
-    # --- 4. THE SPLAT TRICK (OPEN3D) ---
+    # --- 4. OPEN3D COLOR BINDING ---
     pcd = o3d.geometry.PointCloud()
     pcd.points = o3d.utility.Vector3dVector(points)
-    pcd.colors = o3d.utility.Vector3dVector(rgb)
+    pcd.colors = o3d.utility.Vector3dVector(rgb_colors)
 
-    # RE-CENTER: This is the fix for the "Blank Viewer"
-    # It ensures the model is exactly at 0,0,0
+    # 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())
 
-    # DENSITY: Downsample to make points "thicker" and load faster
-    pcd = pcd.voxel_down_sample(voxel_size=0.02) 
+    # Voxelize to make points "solid"
+    pcd = pcd.voxel_down_sample(voxel_size=0.03)
 
-    # --- 5. EXPORT ---
+    # --- 5. SAVE BINARY PLY ---
     temp_dir = tempfile.gettempdir()
     output_path = os.path.join(temp_dir, "model.ply")
-    
-    # write_ascii=False is required for Binary PLY (Colors work best here)
+    # write_ascii=False ensures color data is encoded in a way Gradio understands
     o3d.io.write_point_cloud(output_path, pcd, write_ascii=False)
     
     return output_path, output_path
 
-# --- 6. UI ---
+# --- 6. UI WITH CAMERA FOCUS ---
 with gr.Blocks() as demo:
-    gr.Markdown("## 🪐 3D Splat View (Color-Matched)")
+    gr.Markdown("## 🧊 High-Fidelity 3D Splat View")
+    gr.Markdown("The camera is pre-focused on the object. Click and drag to rotate.")
     
     with gr.Row():
         with gr.Column():
-            img_in = gr.Image(type="pil", label="Upload Photo")
-            btn = gr.Button("🔨 Generate 3D", variant="primary")
+            img_in = gr.Image(type="pil")
+            btn = gr.Button("🔨 Generate & Focus", variant="primary")
         
         with gr.Column():
-            # radius=10 starts the camera at the perfect zoom level
+            # camera_position: (azimuth, elevation, distance)
+            # Distance 12 is perfect for our -5 to 5 coordinate range
             v3d = gr.Model3D(
                 label="3D Viewport",
                 display_mode="solid", 
-                camera_position=(0, 90, 10), 
-                clear_color=(0.08, 0.08, 0.08, 1.0)
+                camera_position=(0, 90, 12), 
+                clear_color=(0.1, 0.1, 0.1, 1.0)
             )
-            dl = gr.DownloadButton("💾 Download .PLY")
+            dl = gr.DownloadButton("💾 Download Colored .PLY")
 
     btn.click(fn=process_to_3d, inputs=[img_in], outputs=[v3d, dl])