Update app.py

app.py

@@ -94,14 +94,7 @@ def allowed_file(filename):
 # Enhanced image preprocessing with better detail preservation
 def preprocess_image(image_path):
     with Image.open(image_path) as img:
-        # Keep alpha channel if present
-        has_alpha = img.mode == 'RGBA'
-        
-        # Convert to proper format while preserving alpha
-        if has_alpha:
-            img = img.convert("RGBA")
-        else:
-            img = img.convert("RGB")
+        img = img.convert("RGB")
         
         # Resize if the image is too large
         if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
@@ -119,17 +112,11 @@ def preprocess_image(image_path):
         # Convert to numpy array for additional preprocessing
         img_array = np.array(img)
         
-        # Extract alpha channel if present
-        if has_alpha:
-            alpha = img_array[:, :, 3]
-            rgb = img_array[:, :, :3]
-        else:
-            rgb = img_array
-            
-        # Apply adaptive histogram equalization for better contrast on RGB channels only
-        if len(rgb.shape) == 3 and rgb.shape[2] == 3:
-            # Convert to LAB color space for better contrast enhancement
-            lab = cv2.cvtColor(rgb, cv2.COLOR_RGB2LAB)
+        # Optional: Apply adaptive histogram equalization for better contrast
+        # This helps the depth model detect more details
+        if len(img_array.shape) == 3 and img_array.shape[2] == 3:
+            # Convert to LAB color space
+            lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
             l, a, b = cv2.split(lab)
             
             # Apply CLAHE to L channel
@@ -140,17 +127,12 @@ def preprocess_image(image_path):
             enhanced_lab = cv2.merge((cl, a, b))
             
             # Convert back to RGB
-            rgb_enhanced = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
+            img_array = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
             
-            # Recombine with alpha if needed
-            if has_alpha:
-                result = np.dstack((rgb_enhanced, alpha))
-                img = Image.fromarray(result, 'RGBA')
-            else:
-                img = Image.fromarray(rgb_enhanced, 'RGB')
+            # Convert back to PIL Image
+            img = Image.fromarray(img_array)
             
         return img
-            
 
 def load_model():
     global depth_estimator, model_loaded, model_loading
@@ -277,7 +259,7 @@ def enhance_depth_map(depth_map, detail_level='medium'):
 
 # Convert depth map to 3D mesh with significantly enhanced detail
 def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
-    """Convert depth map to complete 3D model with all sides"""
+    """Convert depth map to 3D mesh with highly improved detail preservation"""
     # First, enhance the depth map for better details
     enhanced_depth = enhance_depth_map(depth_map, detail_level)
     
@@ -289,94 +271,51 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
     y = np.linspace(0, h-1, resolution)
     x_grid, y_grid = np.meshgrid(x, y)
     
-    # Use bicubic interpolation for smoother surface
+    # Use bicubic interpolation for smoother surface with better details
+    # Create interpolation function
     interp_func = interpolate.RectBivariateSpline(
         np.arange(h), np.arange(w), enhanced_depth, kx=3, ky=3
     )
     
-    # Sample depth at grid points
+    # Sample depth at grid points with the interpolation function
     z_values = interp_func(y, x, grid=True)
     
-    # Process enhancement as in original code
+    # Apply a post-processing step to enhance small details even further
     if detail_level == 'high':
+        # Calculate local gradients to detect edges
         dx = np.gradient(z_values, axis=1) 
         dy = np.gradient(z_values, axis=0)
+        
+        # Enhance edges by increasing depth differences at high gradient areas
         gradient_magnitude = np.sqrt(dx**2 + dy**2)
-        edge_mask = np.clip(gradient_magnitude * 5, 0, 0.2)
+        edge_mask = np.clip(gradient_magnitude * 5, 0, 0.2)  # Scale and limit effect
+        
+        # Apply edge enhancement
         z_values = z_values + edge_mask * (z_values - gaussian_filter(z_values, sigma=1.0))
     
-    # Normalize z-values with advanced scaling
-    z_min, z_max = np.percentile(z_values, [2, 98])
+    # Normalize z-values with advanced scaling for better depth impression
+    z_min, z_max = np.percentile(z_values, [2, 98])  # Remove outliers
     z_values = (z_values - z_min) / (z_max - z_min) if z_max > z_min else z_values
     
-    # Apply depth scaling
+    # Apply depth scaling appropriate to the detail level
     if detail_level == 'high':
-        z_scaling = 2.5
+        z_scaling = 2.5  # More pronounced depth variations
     elif detail_level == 'medium':
-        z_scaling = 2.0
+        z_scaling = 2.0  # Standard depth
     else:
-        z_scaling = 1.5
+        z_scaling = 1.5  # More subtle depth variations
         
     z_values = z_values * z_scaling
     
-    # Normalize coordinates for front face
-    x_grid_front = (x_grid / w - 0.5) * 2.0
-    y_grid_front = (y_grid / h - 0.5) * 2.0
-    
-    # Create all vertices (front, back, and sides)
-    vertices = []
-    
-    # Front face vertices
-    front_vertices = np.vstack([x_grid_front.flatten(), -y_grid_front.flatten(), -z_values.flatten()]).T
-    vertices.append(front_vertices)
-    
-    # Back face vertices (mirrored from front face)
-    back_depth = 1.0  # Constant thickness for the model
-    back_vertices = np.vstack([x_grid_front.flatten(), -y_grid_front.flatten(), -z_values.flatten() - back_depth]).T
-    vertices.append(back_vertices)
-    
-    # Create side vertices (top, bottom, left, right)
-    # For simplicity, we use a grid mapping for sides
-    top_vertices = []
-    bottom_vertices = []
-    left_vertices = []
-    right_vertices = []
-    
-    # Create sides by connecting front and back faces
-    for i in range(resolution):
-        # Top edge
-        for j in range(resolution):
-            if i == 0:
-                top_vertices.append(front_vertices[i * resolution + j])
-                top_vertices.append(back_vertices[i * resolution + j])
-            # Bottom edge
-            if i == resolution - 1:
-                bottom_vertices.append(front_vertices[i * resolution + j])
-                bottom_vertices.append(back_vertices[i * resolution + j])
-            # Left edge
-            if j == 0:
-                left_vertices.append(front_vertices[i * resolution + j])
-                left_vertices.append(back_vertices[i * resolution + j])
-            # Right edge
-            if j == resolution - 1:
-                right_vertices.append(front_vertices[i * resolution + j])
-                right_vertices.append(back_vertices[i * resolution + j])
-    
-    # Combine all vertices
-    all_vertices = np.vstack([
-        front_vertices, 
-        back_vertices, 
-        np.array(top_vertices),
-        np.array(bottom_vertices),
-        np.array(left_vertices),
-        np.array(right_vertices)
-    ])
-    
-    # Create faces (triangles)
-    faces = []
+    # Normalize x and y coordinates
+    x_grid = (x_grid / w - 0.5) * 2.0  # Map to -1 to 1
+    y_grid = (y_grid / h - 0.5) * 2.0  # Map to -1 to 1
     
-    # Front face triangles
-    front_faces = []
+    # Create vertices
+    vertices = np.vstack([x_grid.flatten(), -y_grid.flatten(), -z_values.flatten()]).T
+    
+    # Create faces (triangles) with optimized winding for better normals
+    faces = []
     for i in range(resolution-1):
         for j in range(resolution-1):
             p1 = i * resolution + j
@@ -384,140 +323,97 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
             p3 = (i + 1) * resolution + j
             p4 = (i + 1) * resolution + (j + 1)
             
-            # Calculate normals for consistent orientation
-            v1 = front_vertices[p1]
-            v2 = front_vertices[p2]
-            v3 = front_vertices[p3]
-            v4 = front_vertices[p4]
+            # Calculate normals to ensure consistent orientation
+            v1 = vertices[p1]
+            v2 = vertices[p2]
+            v3 = vertices[p3]
+            v4 = vertices[p4]
             
+            # Calculate normals for both possible triangulations
+            # and choose the one that's more consistent
             norm1 = np.cross(v2-v1, v4-v1)
             norm2 = np.cross(v4-v3, v1-v3)
             
             if np.dot(norm1, norm2) >= 0:
-                front_faces.append([p1, p2, p4])
-                front_faces.append([p1, p4, p3])
+                # Standard triangulation
+                faces.append([p1, p2, p4])
+                faces.append([p1, p4, p3])
             else:
-                front_faces.append([p1, p2, p3])
-                front_faces.append([p2, p4, p3])
+                # Alternative triangulation for smoother surface
+                faces.append([p1, p2, p3])
+                faces.append([p2, p4, p3])
     
-    # Back face triangles (note: reversed winding order for correct normals)
-    back_offset = resolution * resolution  # Offset for back face vertices
-    back_faces = []
-    for i in range(resolution-1):
-        for j in range(resolution-1):
-            p1 = back_offset + i * resolution + j
-            p2 = back_offset + i * resolution + (j + 1)
-            p3 = back_offset + (i + 1) * resolution + j
-            p4 = back_offset + (i + 1) * resolution + (j + 1)
-            
-            # Reverse winding order compared to front face
-            back_faces.append([p1, p4, p2])
-            back_faces.append([p1, p3, p4])
-    
-    # Side faces (connecting front and back)
-    side_faces = []
-    
-    # Add faces for sides (top, bottom, left, right)
-    side_offset = 2 * resolution * resolution  # Offset after front and back
-    
-    # Top side
-    top_count = len(top_vertices)
-    for i in range(0, top_count - 2, 2):
-        side_faces.append([side_offset + i, side_offset + i + 1, side_offset + i + 3])
-        side_faces.append([side_offset + i, side_offset + i + 3, side_offset + i + 2])
-    
-    # Bottom side
-    bottom_offset = side_offset + top_count
-    bottom_count = len(bottom_vertices)
-    for i in range(0, bottom_count - 2, 2):
-        side_faces.append([bottom_offset + i, bottom_offset + i + 3, bottom_offset + i + 1])
-        side_faces.append([bottom_offset + i, bottom_offset + i + 2, bottom_offset + i + 3])
-    
-    # Left side
-    left_offset = bottom_offset + bottom_count
-    left_count = len(left_vertices)
-    for i in range(0, left_count - 2, 2):
-        side_faces.append([left_offset + i, left_offset + i + 1, left_offset + i + 3])
-        side_faces.append([left_offset + i, left_offset + i + 3, left_offset + i + 2])
-    
-    # Right side
-    right_offset = left_offset + left_count
-    right_count = len(right_vertices)
-    for i in range(0, right_count - 2, 2):
-        side_faces.append([right_offset + i, right_offset + i + 3, right_offset + i + 1])
-        side_faces.append([right_offset + i, right_offset + i + 2, right_offset + i + 3])
-    
-    # Combine all faces
-    faces = np.array(front_faces + back_faces + side_faces)
+    faces = np.array(faces)
     
     # Create mesh
-    mesh = trimesh.Trimesh(vertices=all_vertices, faces=faces)
+    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
     
-    # Apply texturing if image is provided
+    # Apply advanced texturing if image is provided
     if image:
-        # Handle RGBA properly to ensure transparency is maintained
-        img_array = np.array(image)
-        
-        # Check if image has alpha channel
-        has_alpha = len(img_array.shape) == 3 and img_array.shape[2] == 4
-        
-        # Create vertex colors with transparency support
-        vertex_colors = np.zeros((all_vertices.shape[0], 4), dtype=np.uint8)
-        
-        # Fill with default color (will be overridden for front face)
-        vertex_colors[:, :3] = [200, 200, 200]  # Light gray default
-        vertex_colors[:, 3] = 255  # Fully opaque
-        
-        # Front face texture (sample from image)
-        for i in range(resolution):
-            for j in range(resolution):
-                # Calculate image coordinates
-                img_x = j * (img_array.shape[1] - 1) / (resolution - 1)
-                img_y = i * (img_array.shape[0] - 1) / (resolution - 1)
-                
-                # Bilinear interpolation setup
-                x0, y0 = int(img_x), int(img_y)
-                x1, y1 = min(x0 + 1, img_array.shape[1] - 1), min(y0 + 1, img_array.shape[0] - 1)
-                
-                # Interpolation weights
-                wx = img_x - x0
-                wy = img_y - y0
-                
-                vertex_idx = i * resolution + j
-                
-                if has_alpha:
-                    # Handle RGBA with bilinear interpolation
-                    for c in range(4):
-                        vertex_colors[vertex_idx, c] = int((1-wx)*(1-wy)*img_array[y0, x0, c] + 
-                                                       wx*(1-wy)*img_array[y0, x1, c] + 
-                                                       (1-wx)*wy*img_array[y1, x0, c] + 
-                                                       wx*wy*img_array[y1, x1, c])
-                else:
-                    # Handle RGB (no alpha)
-                    for c in range(3):
-                        vertex_colors[vertex_idx, c] = int((1-wx)*(1-wy)*img_array[y0, x0, c] + 
-                                                       wx*(1-wy)*img_array[y0, x1, c] + 
-                                                       (1-wx)*wy*img_array[y1, x0, c] + 
-                                                       wx*wy*img_array[y1, x1, c])
-                    vertex_colors[vertex_idx, 3] = 255  # Fully opaque
-        
-        # Apply simpler texturing to back face
-        back_face_start = resolution * resolution
-        back_face_color = [180, 180, 180, 255]  # Slightly darker gray
-        vertex_colors[back_face_start:back_face_start + (resolution * resolution)] = back_face_color
+        # Convert to numpy array if needed
+        if isinstance(image, Image.Image):
+            img_array = np.array(image)
+        else:
+            img_array = image
         
-        mesh.visual.vertex_colors = vertex_colors
+        # Create vertex colors with improved sampling
+        if resolution <= img_array.shape[0] and resolution <= img_array.shape[1]:
+            # Create vertex colors by sampling the image with bilinear interpolation
+            vertex_colors = np.zeros((vertices.shape[0], 4), dtype=np.uint8)
+            
+            # Get normalized coordinates for sampling
+            for i in range(resolution):
+                for j in range(resolution):
+                    # Calculate exact image coordinates with proper scaling
+                    img_x = j * (img_array.shape[1] - 1) / (resolution - 1)
+                    img_y = i * (img_array.shape[0] - 1) / (resolution - 1)
+                    
+                    # Bilinear interpolation for smooth color transitions
+                    x0, y0 = int(img_x), int(img_y)
+                    x1, y1 = min(x0 + 1, img_array.shape[1] - 1), min(y0 + 1, img_array.shape[0] - 1)
+                    
+                    # Calculate interpolation weights
+                    wx = img_x - x0
+                    wy = img_y - y0
+                    
+                    vertex_idx = i * resolution + j
+                    
+                    if len(img_array.shape) == 3 and img_array.shape[2] == 3:  # RGB
+                        # Perform bilinear interpolation for each color channel
+                        r = int((1-wx)*(1-wy)*img_array[y0, x0, 0] + wx*(1-wy)*img_array[y0, x1, 0] + 
+                                (1-wx)*wy*img_array[y1, x0, 0] + wx*wy*img_array[y1, x1, 0])
+                        g = int((1-wx)*(1-wy)*img_array[y0, x0, 1] + wx*(1-wy)*img_array[y0, x1, 1] + 
+                                (1-wx)*wy*img_array[y1, x0, 1] + wx*wy*img_array[y1, x1, 1])
+                        b = int((1-wx)*(1-wy)*img_array[y0, x0, 2] + wx*(1-wy)*img_array[y0, x1, 2] + 
+                                (1-wx)*wy*img_array[y1, x0, 2] + wx*wy*img_array[y1, x1, 2])
+                        
+                        vertex_colors[vertex_idx, :3] = [r, g, b]
+                        vertex_colors[vertex_idx, 3] = 255  # Alpha
+                    elif len(img_array.shape) == 3 and img_array.shape[2] == 4:  # RGBA
+                        for c in range(4):  # For each RGBA channel
+                            vertex_colors[vertex_idx, c] = int((1-wx)*(1-wy)*img_array[y0, x0, c] + 
+                                                            wx*(1-wy)*img_array[y0, x1, c] + 
+                                                            (1-wx)*wy*img_array[y1, x0, c] + 
+                                                            wx*wy*img_array[y1, x1, c])
+                    else:
+                        # Handle grayscale with bilinear interpolation
+                        gray = int((1-wx)*(1-wy)*img_array[y0, x0] + wx*(1-wy)*img_array[y0, x1] + 
+                                  (1-wx)*wy*img_array[y1, x0] + wx*wy*img_array[y1, x1])
+                        vertex_colors[vertex_idx, :3] = [gray, gray, gray]
+                        vertex_colors[vertex_idx, 3] = 255
+            
+            mesh.visual.vertex_colors = vertex_colors
     
     # Apply smoothing to get rid of staircase artifacts
     if detail_level != 'high':
+        # For medium and low detail, apply Laplacian smoothing
+        # but preserve the overall shape
         mesh = mesh.smoothed(method='laplacian', iterations=1)
     
     # Calculate and fix normals for better rendering
     mesh.fix_normals()
     
     return mesh
-    
-        
 
 @app.route('/health', methods=['GET'])
 def health_check():