Update app.py

app.py

@@ -14,12 +14,17 @@ midas_transform = torch.hub.load("intel-isl/MiDaS", "transforms").default_transf
 
 def estimate_depth(image):
     """Estimate depth map using MiDaS v3."""
-    image = image.convert("RGB")
-    image = np.array(image) / 255.0  # Normalize and convert to NumPy array
-    input_tensor = midas_transform(image).unsqueeze(0).to(device)
+    image = image.convert("RGB")  # Ensure it's in RGB format
+    img_tensor = midas_transform(image).unsqueeze(0).to(device)
+
+    # Ensure tensor shape is [1, 3, H, W]
+    if img_tensor.dim() == 5:  # If an extra batch dimension is present
+        img_tensor = img_tensor.squeeze(1)
+
     with torch.no_grad():
-        depth = midas_model(input_tensor).squeeze().cpu().numpy()
-    depth = cv2.resize(depth, (image.shape[1], image.shape[0]))
+        depth = midas_model(img_tensor).squeeze().cpu().numpy()
+    
+    depth = cv2.resize(depth, (image.size[0], image.size[1]))
     depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255
     return depth.astype(np.uint8)