Update app.py

app.py

@@ -51,17 +51,17 @@ def detect_and_annotate(
     masks = None
     if results.masks is not None:
         masks = results.masks.data.cpu().numpy()
+        # Reshape masks to (num_masks, H, W)
+        masks = np.transpose(masks, (1, 2, 0))  # From (H, W, num_masks) to (num_masks, H, W)
+        
         # Resize masks to match original image dimensions
         h, w = image.shape[:2]
-        masks = [
-            cv2.resize(mask.astype(float), (w, h), 
-            interpolation=cv2.INTER_LINEAR
-        ).astype(bool)
-            for mask in masks
-        ]
-        masks = np.array(masks)
-        # Transpose to (H, W, num_masks)
-        masks = np.transpose(masks, (1, 2, 0))
+        resized_masks = []
+        for mask in masks:
+            resized_mask = cv2.resize(mask.astype(float), (w, h), interpolation=cv2.INTER_LINEAR)
+            resized_masks.append(resized_mask)
+        masks = np.array(resized_masks)
+        masks = masks.astype(bool)
     
     # Create Detections object
     detections = sv.Detections(