Update app.py

app.py

@@ -3,22 +3,30 @@ import gradio as gr
 from PIL import Image
 import numpy as np
 import open3d as o3d
-from transformers import pipeline
+import cv2
 
-class CPUFriendlyImageTo3DConverter:
+class RobustImageTo3DConverter:
     def __init__(self):
-        # Use lighter depth estimation model
-        self.depth_estimator = pipeline(
-            "depth-estimation", 
-            model="facebook/dpt-small"  # Lighter model for CPU
-        )
+        # Use OpenCV for depth estimation
+        self.use_midas = False
     
     def estimate_depth(self, image):
         """
-        Estimate depth using a lightweight model
+        Estimate depth using OpenCV's stereo algorithms
         """
-        depth_result = self.depth_estimator(image)
-        depth_map = np.array(depth_result['depth'])
+        # Convert PIL Image to OpenCV format
+        img_array = np.array(image)
+        
+        # Convert to grayscale
+        gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
+        
+        # Create simple depth estimation using edge detection and blurring
+        edges = cv2.Canny(gray, 100, 200)
+        depth_map = cv2.distanceTransform(255 - edges, cv2.DIST_L2, 5)
+        
+        # Normalize depth map
+        depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX)
+        
         return depth_map
     
     def create_point_cloud(self, image, depth_map):
@@ -31,13 +39,17 @@ class CPUFriendlyImageTo3DConverter:
         # Downsample to reduce computational load
         step = 5  # Reduce resolution
         points = []
+        colors = []
+        
         for y in range(0, height, step):
             for x in range(0, width, step):
-                z = depth_map[y, x]
-                points.append([x, y, z])
+                z = depth_map[y, x] / 255.0  # Normalize depth
+                points.append([x, y, z * 10])  # Scale depth
+                colors.append(img_array[y, x] / 255.0)  # Normalize colors
         
         pcd = o3d.geometry.PointCloud()
         pcd.points = o3d.utility.Vector3dVector(points)
+        pcd.colors = o3d.utility.Vector3dVector(colors)
         
         return pcd
     
@@ -46,12 +58,16 @@ class CPUFriendlyImageTo3DConverter:
         Simplified mesh reconstruction for CPU
         """
         point_cloud.estimate_normals()
+        point_cloud.orient_normals_consistent_tangent_plane(100)
         
         # Use simpler mesh reconstruction with lower depth
         mesh, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
             point_cloud, depth=6  # Reduced depth for faster processing
         )
         
+        # Color the mesh
+        mesh.paint_uniform_color([0.7, 0.7, 0.7])
+        
         return mesh
     
     def process_image(self, input_image):
@@ -74,7 +90,7 @@ class CPUFriendlyImageTo3DConverter:
         return output_path
 
 def create_huggingface_space():
-    converter = CPUFriendlyImageTo3DConverter()
+    converter = RobustImageTo3DConverter()
     
     def convert_image(input_image):
         try:
@@ -82,6 +98,10 @@ def create_huggingface_space():
             if not isinstance(input_image, Image.Image):
                 input_image = Image.fromarray(input_image)
             
+            # Resize image if too large
+            max_size = (800, 800)
+            input_image.thumbnail(max_size, Image.LANCZOS)
+            
             # Process image
             output_model = converter.process_image(input_image)
             return output_model
@@ -96,8 +116,8 @@ def create_huggingface_space():
             gr.File(label="3D Model (OBJ)"),
             gr.Textbox(label="Conversion Status")
         ],
-        title="CPU-Friendly Image to 3D Model Converter",
-        description="Convert images to 3D models using lightweight depth estimation and point cloud reconstruction."
+        title="Robust Image to 3D Model Converter",
+        description="Convert images to 3D models using OpenCV depth estimation and point cloud reconstruction."
     )
     
     return iface