Update app.py

app.py

@@ -12,6 +12,24 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
 processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
 
+def process_mask(mask, target_size):
+    # Ensure mask is 2D
+    if mask.ndim > 2:
+        mask = mask.squeeze()
+    
+    # If mask is still not 2D, take the first 2D slice
+    if mask.ndim > 2:
+        mask = mask[0]
+    
+    # Convert to binary
+    mask = (mask > 0.5).astype(np.uint8) * 255
+    
+    # Resize mask to match original image size using PIL
+    mask_image = Image.fromarray(mask)
+    mask_image = mask_image.resize(target_size, Image.NEAREST)
+    
+    return np.array(mask_image) > 0
+
 def segment_image(input_image, segment_anything):
     try:
         if input_image is None:
@@ -44,20 +62,13 @@ def segment_image(input_image, segment_anything):
             inputs["reshaped_input_sizes"].cpu()
         )
         
-        # Convert mask to numpy array
+        # Process the mask
         if segment_anything:
             combined_mask = np.any(masks[0].numpy() > 0.5, axis=0)
         else:
-            combined_mask = masks[0][0].numpy() > 0.5
-        
-        # Ensure mask is 2D
-        if combined_mask.ndim > 2:
-            combined_mask = combined_mask.squeeze()
+            combined_mask = masks[0][0].numpy()
         
-        # Resize mask to match original image size using PIL
-        mask_image = Image.fromarray((combined_mask * 255).astype(np.uint8))
-        mask_image = mask_image.resize(original_size, Image.NEAREST)
-        combined_mask = np.array(mask_image) > 0
+        combined_mask = process_mask(combined_mask, original_size)
         
         # Overlay the mask on the original image
         result_image = np.array(input_image)