Update app.py

app.py

@@ -20,11 +20,14 @@ global_state = {
 
 # Helper to apply mask overlay
 def apply_mask_overlay(image: Image.Image, mask: np.ndarray, color=(255, 0, 0)) -> Image.Image:
-    mask_img = Image.fromarray(mask.astype(np.uint8) * 255).convert("L")
+    if mask.ndim == 3:
+        mask = mask.squeeze()
+    if mask.max() <= 1:
+        mask = (mask * 255).astype(np.uint8)
+    mask_img = Image.fromarray(mask).convert("L")
     color_mask = Image.new("RGB", image.size, color)
-    mask_rgb = Image.composite(color_mask, image, mask_img)
-    blended = Image.blend(image, mask_rgb, alpha=0.5)
-    return blended
+    blended = Image.composite(color_mask, image, mask_img)
+    return Image.blend(image, blended, alpha=0.5)
 
 # Set image
 def upload_image(img):
@@ -56,9 +59,15 @@ def run_segmentation():
     inputs = processor(image, return_tensors="pt").to(device)
 
     if global_state["clicks"]:
-        points = torch.tensor([[[x, y] for (x, y, l) in global_state["clicks"]]], device=device)
-        labels = torch.tensor([[l for (_, _, l) in global_state["clicks"]]], device=device)
-        inputs.update({"input_points": points, "input_labels": labels})
+        coords = [[(x, y) for (x, y, l) in global_state["clicks"]]]
+        labels = [[l for (_, _, l) in global_state["clicks"]]]
+        input_points = torch.tensor([coords], device=device)  # shape [1, 1, N, 2]
+        input_labels = torch.tensor([labels], device=device)  # shape [1, 1, N]
+
+        inputs.update({
+            "input_points": input_points,
+            "input_labels": input_labels
+            })
 
     if global_state["bbox"]:
         x0, y0, x1, y1 = global_state["bbox"]
@@ -74,7 +83,7 @@ def run_segmentation():
         inputs["reshaped_input_sizes"].cpu()
     )[0]
 
-    final_mask = masks[0].numpy()
+    final_mask = masks[0].numpy().astype(np.uint8)  # shape: (H, W)
     overlayed = apply_mask_overlay(image.convert("RGB"), final_mask)
 
     return overlayed, "Segmentation complete."