Update app.py

app.py

@@ -25,6 +25,9 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
 
     pred = torch.sigmoid(preds)
     mat = pred.cpu().numpy()
+
+    # Ensure we are working with a single-channel 2D mask
+    mat = np.squeeze(mat, axis=0)  # Remove batch dimension if it exists
     mask = Image.fromarray(np.uint8(mat * 255), "L")
     mask = mask.convert("RGB")
     mask = mask.resize(image.size)
@@ -37,7 +40,6 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
 
     # threshold the mask
     bmask = mask > threshold
-    # zero out values below the threshold
     mask[mask < threshold] = 0
 
     fig, ax = plt.subplots()
@@ -74,6 +76,7 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
 
     return fig, result_mask, result_output
 
+
     # Existing process_image function, copy it here
     # ...