Update app.py

app.py

@@ -25,6 +25,10 @@ def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
 
     pred = torch.sigmoid(preds)
     mat = pred.cpu().numpy()
+
+    # Squeeze to remove dimensions of size 1 (batch and channel dimensions)
+    mat = np.squeeze(mat, axis=(0, 1))  # Convert from [1, 1, H, W] to [H, W]
+
     mask = Image.fromarray(np.uint8(mat * 255), "L")
     mask = mask.convert("RGB")
     mask = mask.resize(image.size)