Update app.py

app.py

@@ -259,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 3D mesh with highly improved detail preservation"""
+    """Convert depth map to 3D mesh with highly improved detail preservation and transparency support"""
     # First, enhance the depth map for better details
     enhanced_depth = enhance_depth_map(depth_map, detail_level)
     
@@ -314,7 +314,38 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
     # Create vertices
     vertices = np.vstack([x_grid.flatten(), -y_grid.flatten(), -z_values.flatten()]).T
     
-    # Create faces (triangles) with optimized winding for better normals
+    # Create transparency mask for the image if it has an alpha channel
+    # This will be used to filter out faces that contain transparent pixels
+    has_alpha = False
+    alpha_mask = np.ones((resolution, resolution), dtype=bool)
+    
+    if image is not None:
+        if isinstance(image, Image.Image):
+            if image.mode == 'RGBA':
+                has_alpha = True
+                # Convert image to numpy array with alpha channel
+                img_array = np.array(image)
+                # Extract alpha channel
+                alpha_channel = img_array[:, :, 3]
+                
+                # Create alpha mask by sampling alpha channel at grid positions
+                for i in range(resolution):
+                    for j in range(resolution):
+                        img_y = int(i * (img_array.shape[0] - 1) / (resolution - 1))
+                        img_x = int(j * (img_array.shape[1] - 1) / (resolution - 1))
+                        alpha_mask[i, j] = alpha_channel[img_y, img_x] > 10  # Threshold for transparency
+        elif isinstance(image, np.ndarray) and image.shape[2] == 4:  # RGBA numpy array
+            has_alpha = True
+            alpha_channel = image[:, :, 3]
+            
+            # Sample alpha channel at grid positions
+            for i in range(resolution):
+                for j in range(resolution):
+                    img_y = int(i * (image.shape[0] - 1) / (resolution - 1))
+                    img_x = int(j * (image.shape[1] - 1) / (resolution - 1))
+                    alpha_mask[i, j] = alpha_channel[img_y, img_x] > 10  # Threshold for transparency
+    
+    # Create faces (triangles) with transparency handling
     faces = []
     for i in range(resolution-1):
         for j in range(resolution-1):
@@ -323,25 +354,31 @@ 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 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:
-                # Standard triangulation
-                faces.append([p1, p2, p4])
-                faces.append([p1, p4, p3])
-            else:
-                # Alternative triangulation for smoother surface
-                faces.append([p1, p2, p3])
-                faces.append([p2, p4, p3])
+            # Only create faces if all vertices are visible (non-transparent)
+            if not has_alpha or (alpha_mask[i, j] and alpha_mask[i, j+1] and 
+                                alpha_mask[i+1, j] and alpha_mask[i+1, j+1]):
+                # 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:
+                    # Standard triangulation
+                    faces.append([p1, p2, p4])
+                    faces.append([p1, p4, p3])
+                else:
+                    # Alternative triangulation for smoother surface
+                    faces.append([p1, p2, p3])
+                    faces.append([p2, p4, p3])
+    
+    if len(faces) == 0:
+        raise ValueError("No faces generated - image may be completely transparent")
     
     faces = np.array(faces)
     
@@ -356,10 +393,11 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
         else:
             img_array = image
         
-        # Create vertex colors with improved sampling
+        # Create vertex colors with improved sampling and transparency support
         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)
+            vertex_colors[:, 3] = 255  # Default alpha to opaque
             
             # Get normalized coordinates for sampling
             for i in range(resolution):
@@ -378,23 +416,26 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
                     
                     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])
+                    # Apply vertex colors based on image format
+                    if len(img_array.shape) == 3:
+                        if img_array.shape[2] == 4:  # RGBA
+                            # Set colors with alpha channel
+                            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])
+                        elif img_array.shape[2] == 3:  # RGB
+                            # Apply 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  # Full opacity for RGB images
                     else:
                         # Handle grayscale with bilinear interpolation
                         gray = int((1-wx)*(1-wy)*img_array[y0, x0] + wx*(1-wy)*img_array[y0, x1] + 
@@ -415,6 +456,60 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
     
     return mesh
 
+    # Enhanced image preprocessing to properly handle PNGs with transparency
+def preprocess_image(image_path):
+    with Image.open(image_path) as img:
+        # Get original mode to check if it has transparency
+        original_mode = img.mode
+        
+        # Convert to RGB or RGBA as needed
+        if original_mode == 'RGBA':
+            # Keep alpha channel for transparency
+            img = img.convert("RGBA")
+        else:
+            # Otherwise use RGB
+            img = img.convert("RGB")
+        
+        # Resize if the image is too large
+        if img.width > MAX_DIMENSION or img.height > MAX_DIMENSION:
+            # Calculate new dimensions while preserving aspect ratio
+            if img.width > img.height:
+                new_width = MAX_DIMENSION
+                new_height = int(img.height * (MAX_DIMENSION / img.width))
+            else:
+                new_height = MAX_DIMENSION
+                new_width = int(img.width * (MAX_DIMENSION / img.height))
+            
+            # Use high-quality Lanczos resampling for better detail preservation
+            img = img.resize((new_width, new_height), Image.LANCZOS)
+        
+        # Handle enhancement only for RGB (non-transparent) parts
+        if original_mode != 'RGBA':
+            # Convert to numpy array for additional preprocessing
+            img_array = np.array(img)
+            
+            # 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
+                clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+                cl = clahe.apply(l)
+                
+                # Merge channels back
+                enhanced_lab = cv2.merge((cl, a, b))
+                
+                # Convert back to RGB
+                img_array = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
+                
+                # Convert back to PIL Image
+                img = Image.fromarray(img_array)
+        
+        return img
+
 @app.route('/health', methods=['GET'])
 def health_check():
     return jsonify({