added multi-segmentation support

app.py

@@ -105,26 +105,37 @@ def detection_fn(image, prompt):
 def segmentation_fn(image, prompt):
     image_np = np.array(image)
     image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    
+    # Detect with Grounding DINO
     detections, _ = dino_model.predict_with_caption(
         image=image_cv, caption=prompt, box_threshold=0.35, text_threshold=0.25
     )
+    
     boxes = detections.xyxy
     sam_predictor.set_image(image_np)
-    masks, scores, _ = sam_predictor.predict(box=boxes, multimask_output=True)
-    if masks is None or len(masks) == 0:
+    
+    all_masks = []
+    for box in boxes:
+        box = box.reshape(1, 4)
+        masks, scores, _ = sam_predictor.predict(box=box, multimask_output=True)
+        if masks is not None:
+            all_masks.append(masks[np.argmax(scores)])
+    
+    if not all_masks:
         raise ValueError("No masks found")
-    mask = masks[np.argmax(scores)]
+    
+    # Combine masks into one binary mask
+    merged_mask = np.any(all_masks, axis=0).astype(np.uint8) * 255
 
-    # Visualize mask
+    # Overlay on image
     def overlay_mask(mask, image):
-        color = np.concatenate([np.random.random(3), np.array([0.8])])
-        h, w = mask.shape[-2:]
-        mask_img = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
-        image_pil = Image.fromarray(image).convert("RGBA")
-        mask_pil = Image.fromarray((mask_img * 255).astype(np.uint8)).convert("RGBA")
-        return np.array(Image.alpha_composite(image_pil, mask_pil))
-
-    return overlay_mask(mask, image_np)
+        color = np.array([0, 255, 0], dtype=np.uint8)  # Green
+        mask_rgb = np.stack([mask] * 3, axis=-1)
+        overlay = np.where(mask_rgb, color, image)
+        return overlay
+
+    return overlay_mask(merged_mask, image_np)
+
 
 def inpainting_fn(image, prompt):
     image_np = np.array(image)