Update app.py

app.py

@@ -51,7 +51,7 @@ def plot_depth_map(depth_map, original_image):
     
     # Resize original image to match depth map dimensions
     original_image_resized = cv2.resize(original_image, (depth_map.shape[1], depth_map.shape[0]))
-    colors = original_image_resized.reshape(-1, 3) / 255.0
+    colors = original_image_resized.reshape(depth_map.shape[0], depth_map.shape[1], 3) / 255.0
     
     ax.plot_surface(x, y, depth_map, facecolors=colors, shade=False)
     ax.view_init(elev=70, azim=0)  # Tilt by 20 degrees (90 - 20 = 70)
@@ -75,7 +75,11 @@ def process_frame(image):
     
     depth_map = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min())
     
-    return plot_depth_map(depth_map, image)  # Pass the original image, not resized
+    # Convert BGR to RGB if necessary
+    if image.shape[2] == 3:  # Check if it's a color image
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    
+    return plot_depth_map(depth_map, image)
 
 interface = gr.Interface(
     fn=process_frame,