Update app.py

app.py

@@ -30,21 +30,23 @@ def transcribe(audio):
 def segment_image(image: Image.Image):
     results = segmentation_pipeline(image)
 
-    # Generate a random color for each object
+    # Make a copy for overlay
     overlay = np.array(image).copy()
     annotations = []
+
     for r in results:
-        mask = np.array(r["mask"])  # mask is a PIL image
+        mask = np.array(r["mask"]) > 0  # ensure binary mask
         label = r["label"]
 
-        # Random color per mask
+        # Random color for each object
         color = [random.randint(0, 255) for _ in range(3)]
+        overlay[mask] = (0.6 * overlay[mask] + 0.4 * np.array(color)).astype(np.uint8)
 
-        # Apply semi-transparent overlay
-        overlay[mask > 0] = (0.6 * overlay[mask > 0] + 0.4 * np.array(color)).astype(np.uint8)
+        # Append as (mask, label) where mask is np.ndarray
+        annotations.append((mask, label))
 
-        # Store mask + label for Gradio
-        annotations.append((r["mask"], label))
+    overlay_img = Image.fromarray(overlay)
+    return (overlay_img, annotations)
 
     overlay_img = Image.fromarray(overlay)
     return (overlay_img, annotations)