Update app.py

app.py

@@ -139,10 +139,11 @@ def preprocess_image(image_path):
                 new_width = int(img.width * (MAX_DIMENSION / img.height))
             img = img.resize((new_width, new_height), Image.LANCZOS)
         
-        # Remove background
-        img = remove_background(img)
+        # Remove background and convert back to RGB for processor
+        img_with_alpha = remove_background(img)
+        img_rgb = img_with_alpha.convert("RGB")  # Convert to RGB for processor
         
-        img_array = np.array(img)
+        img_array = np.array(img_rgb)
         if len(img_array.shape) == 3 and img_array.shape[2] == 3:
             lab = cv2.cvtColor(img_array, cv2.COLOR_RGB2LAB)
             l, a, b = cv2.split(lab)
@@ -150,9 +151,9 @@ def preprocess_image(image_path):
             cl = clahe.apply(l)
             enhanced_lab = cv2.merge((cl, a, b))
             img_array = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2RGB)
-            img = Image.fromarray(img_array)
+            img_rgb = Image.fromarray(img_array)
         
-        return img
+        return img_rgb  # Return RGB image
 
 def load_models():
     global dpt_estimator, depth_anything_model, depth_anything_processor, model_loaded, model_loading
@@ -293,7 +294,7 @@ def enhance_depth_map(depth_map, detail_level='medium'):
         smooth1 = gaussian_filter(enhanced_depth, sigma=0.3)
         smooth2 = gaussian_filter(enhanced_depth, sigma=1.5)
         edge_mask = enhanced_depth - smooth2
-        enhanced_depth = smooth1 + 0.8 * edge_mask  # Reduced enhancement
+        enhanced_depth = smooth1 + 0.8 * edge_mask
     elif detail_level == 'medium':
         blurred = gaussian_filter(enhanced_depth, sigma=0.7)
         mask = enhanced_depth - blurred
@@ -357,7 +358,7 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
     mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
     
     if image:
-        img_array = np.array(image)
+        img_array = np.array(image.convert("RGB"))  # Ensure RGB for consistency
         vertex_colors = np.zeros((vertices.shape[0], 4), dtype=np.uint8)
         for i in range(resolution):
             for j in range(resolution):
@@ -368,15 +369,21 @@ def depth_to_mesh(depth_map, image, resolution=100, detail_level='medium'):
                 wx = img_x - x0
                 wy = img_y - y0
                 vertex_idx = i * resolution + j
-                if len(img_array.shape) == 3 and img_array.shape[2] == 4:
-                    for c in range(4):
+                if len(img_array.shape) == 3 and img_array.shape[2] == 3:
+                    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
+                elif len(img_array.shape) == 3 and img_array.shape[2] == 4:
+                    for c in range(3):  # Use only RGB channels
                         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:
-                    r, g, b = img_array[y0, x0]
-                    vertex_colors[vertex_idx, :3] = [r, g, b]
                     vertex_colors[vertex_idx, 3] = 255
         mesh.visual.vertex_colors = vertex_colors