Update app.py

app.py

@@ -16,6 +16,7 @@ from flask_cors import CORS
 import numpy as np
 import trimesh
 from transformers import pipeline
+from scipy.ndimage import gaussian_filter, uniform_filter
 
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
@@ -48,7 +49,7 @@ model_loaded = False
 model_loading = False
 
 # Configuration for processing
-TIMEOUT_SECONDS = 180  # 3 minutes max for processing
+TIMEOUT_SECONDS = 240  # 4 minutes max for processing (increased for larger model)
 MAX_DIMENSION = 512    # Max image dimension to process
 
 # TimeoutError for handling timeouts
@@ -121,8 +122,8 @@ def load_model():
         model_loading = True
         print("Starting model loading...")
         
-        # Using DPT-Hybrid which is smaller than other depth estimation models
-        model_name = "Intel/dpt-hybrid-midas"
+        # Using DPT-Large which provides better detail than DPT-Hybrid
+        model_name = "Intel/dpt-large"
         
         # Download model with retry mechanism
         max_retries = 3
@@ -170,20 +171,23 @@ def load_model():
     finally:
         model_loading = False
 
-# Convert depth map to 3D mesh
+# Convert depth map to 3D mesh with enhanced detail
 def depth_to_mesh(depth_map, image, resolution=100):
-    """Convert depth map to 3D mesh"""
+    """Convert depth map to 3D mesh with improved detail preservation"""
     # Convert depth_map to numpy array if it's a PIL Image
     if isinstance(depth_map, Image.Image):
         depth_map = np.array(depth_map)
     
     # Make sure the depth map is 2D
     if len(depth_map.shape) > 2:
-        # If it's a 3D array (like RGB), convert to grayscale
         depth_map = np.mean(depth_map, axis=2) if depth_map.shape[2] > 1 else depth_map[:,:,0]
     
+    # Apply bilateral filter to smooth the depth map while preserving edges
+    # First, apply a slight gaussian filter to remove noise
+    depth_map_smooth = gaussian_filter(depth_map, sigma=1.0)
+    
     # Get dimensions
-    h, w = depth_map.shape
+    h, w = depth_map_smooth.shape
     
     # Create a grid of points
     x = np.linspace(0, w-1, resolution)
@@ -193,11 +197,24 @@ def depth_to_mesh(depth_map, image, resolution=100):
     # Sample depth at grid points
     x_indices = x_grid.astype(int)
     y_indices = y_grid.astype(int)
-    z_values = depth_map[y_indices, x_indices]
-    
-    # Normalize depth values to suitable range
-    z_min, z_max = z_values.min(), z_values.max()
-    z_values = (z_values - z_min) / (z_max - z_min) * 2.0  # Map to 0-2 range
+    z_values = depth_map_smooth[y_indices, x_indices]
+    
+    # Normalize depth values with better scaling
+    z_min, z_max = np.percentile(z_values, [2, 98])  # Removes outliers
+    z_values = (z_values - z_min) / (z_max - z_min) if z_max > z_min else z_values
+    z_values = z_values * 2.0  # Scale depth
+    
+    # Apply a local contrast enhancement to bring out details
+    # Simple adaptive normalization
+    window_size = resolution // 10
+    if window_size > 0:
+        local_mean = uniform_filter(z_values, size=window_size)
+        local_var = uniform_filter(z_values**2, size=window_size) - local_mean**2
+        local_std = np.sqrt(np.maximum(local_var, 0))
+        
+        # Enhance local contrast
+        enhanced_z = (z_values - local_mean) / (local_std + 0.01) * 0.5 + z_values
+        z_values = np.clip(enhanced_z, 0, None)  # Keep values positive
     
     # Normalize x and y coordinates
     x_grid = (x_grid / w - 0.5) * 2.0  # Map to -1 to 1
@@ -215,6 +232,7 @@ def depth_to_mesh(depth_map, image, resolution=100):
             p3 = (i + 1) * resolution + j
             p4 = (i + 1) * resolution + (j + 1)
             
+            # Create two triangles for each grid cell
             faces.append([p1, p2, p4])
             faces.append([p1, p4, p3])
     
@@ -223,10 +241,37 @@ def depth_to_mesh(depth_map, image, resolution=100):
     # Create mesh
     mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
     
-    # Optional: Apply texture from original image
+    # Apply texturing if image is provided
     if image:
-        # This is simplified - proper UV mapping would be needed for accurate texturing
-        pass
+        # Convert to numpy array if needed
+        if isinstance(image, Image.Image):
+            img_array = np.array(image)
+        else:
+            img_array = image
+            
+        # Create simple texture by sampling the original image
+        if resolution <= img_array.shape[0] and resolution <= img_array.shape[1]:
+            # Create vertex colors by sampling the image
+            vertex_colors = np.zeros((vertices.shape[0], 4), dtype=np.uint8)
+            
+            for i in range(resolution):
+                for j in range(resolution):
+                    img_x = min(int(j * img_array.shape[1] / resolution), img_array.shape[1]-1)
+                    img_y = min(int(i * img_array.shape[0] / resolution), img_array.shape[0]-1)
+                    
+                    vertex_idx = i * resolution + j
+                    if len(img_array.shape) == 3 and img_array.shape[2] == 3:  # RGB
+                        vertex_colors[vertex_idx, :3] = img_array[img_y, img_x, :]
+                        vertex_colors[vertex_idx, 3] = 255  # Alpha
+                    elif len(img_array.shape) == 3 and img_array.shape[2] == 4:  # RGBA
+                        vertex_colors[vertex_idx, :] = img_array[img_y, img_x, :]
+                    else:
+                        # Handle grayscale or other formats
+                        gray_value = img_array[img_y, img_x]
+                        vertex_colors[vertex_idx, :3] = [gray_value, gray_value, gray_value]
+                        vertex_colors[vertex_idx, 3] = 255
+            
+            mesh.visual.vertex_colors = vertex_colors
     
     return mesh
 
@@ -234,7 +279,7 @@ def depth_to_mesh(depth_map, image, resolution=100):
 def health_check():
     return jsonify({
         "status": "healthy", 
-        "model": "Depth-Based 3D Model Generator",
+        "model": "Depth-Based 3D Model Generator (DPT-Large)",
         "device": "cuda" if torch.cuda.is_available() else "cpu"
     }), 200
 
@@ -294,6 +339,7 @@ def convert_image_to_3d():
     try:
         mesh_resolution = min(int(request.form.get('mesh_resolution', 100)), 200)  # Limit max resolution
         output_format = request.form.get('output_format', 'obj').lower()
+        detail_level = request.form.get('detail_level', 'medium').lower()  # New parameter for detail level
     except ValueError:
         return jsonify({"error": "Invalid parameter values"}), 400
     
@@ -301,6 +347,12 @@ def convert_image_to_3d():
     if output_format not in ['obj', 'glb']:
         return jsonify({"error": "Unsupported output format. Use 'obj' or 'glb'"}), 400
     
+    # Adjust mesh resolution based on detail level
+    if detail_level == 'high':
+        mesh_resolution = min(mesh_resolution * 1.5, 200)
+    elif detail_level == 'low':
+        mesh_resolution = max(int(mesh_resolution * 0.7), 50)
+    
     # Create a job ID
     job_id = str(uuid.uuid4())
     output_dir = os.path.join(RESULTS_FOLDER, job_id)
@@ -372,7 +424,8 @@ def convert_image_to_3d():
                 processing_jobs[job_id]['progress'] = 60
                 
                 # Create mesh from depth map
-                mesh = depth_to_mesh(depth_map, image, resolution=mesh_resolution)
+                mesh_resolution_int = int(mesh_resolution)
+                mesh = depth_to_mesh(depth_map, image, resolution=mesh_resolution_int)
                 processing_jobs[job_id]['progress'] = 80
                 
             except Exception as e:
@@ -523,8 +576,13 @@ def cleanup_old_jobs():
 @app.route('/', methods=['GET'])
 def index():
     return jsonify({
-        "message": "Image to 3D API is running", 
-        "endpoints": ["/convert", "/progress/<job_id>", "/download/<job_id>", "/preview/<job_id>"]
+        "message": "Image to 3D API is running (DPT-Large Model)", 
+        "endpoints": ["/convert", "/progress/<job_id>", "/download/<job_id>", "/preview/<job_id>"],
+        "parameters": {
+            "mesh_resolution": "Integer (50-200), controls mesh density",
+            "output_format": "obj or glb",
+            "detail_level": "low, medium, or high"
+        }
     }), 200
 
 if __name__ == '__main__':