Update app.py

app.py

@@ -3,8 +3,8 @@ import gradio as gr
 import torch
 import numpy as np
 from PIL import Image
-import matplotlib
 from transformers import Sam3Processor, Sam3Model
+import requests
 import warnings
 warnings.filterwarnings("ignore")
 
@@ -13,44 +13,21 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 model = Sam3Model.from_pretrained("facebook/sam3", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32).to(device)
 processor = Sam3Processor.from_pretrained("facebook/sam3")
 
-def overlay_masks(image: Image.Image, masks: torch.Tensor) -> Image.Image:
-    """
-    Overlay segmentation masks on the input image using rainbow colormap.
-    """
-    image = image.convert("RGBA")
-    masks = 255 * masks.cpu().numpy().astype(np.uint8)
-    
-    n_masks = masks.shape[0]
-    if n_masks == 0:
-        return image.convert("RGB")
-    
-    cmap = matplotlib.colormaps.get_cmap("rainbow").resampled(n_masks)
-    colors = [
-        tuple(int(c * 255) for c in cmap(i)[:3])
-        for i in range(n_masks)
-    ]
-
-    for mask, color in zip(masks, colors):
-        mask_img = Image.fromarray(mask)
-        overlay = Image.new("RGBA", image.size, color + (0,))
-        alpha = mask_img.point(lambda v: int(v * 0.5))
-        overlay.putalpha(alpha)
-        image = Image.alpha_composite(image, overlay)
-    return image
-
 @spaces.GPU()
 def segment(image: Image.Image, text: str, threshold: float, mask_threshold: float):
     """
     Perform promptable concept segmentation using SAM3.
+    Returns format compatible with gr.AnnotatedImage: (image, [(mask, label), ...])
     """
     if image is None:
         return None, "❌ Please upload an image."
     
+    if not text.strip():
+        return (image, []), "❌ Please enter a text prompt."
+    
     try:
-        # Ensure inputs match model dtype
         inputs = processor(images=image, text=text.strip(), return_tensors="pt").to(device)
         
-        # Convert inputs to match model dtype
         for key in inputs:
             if inputs[key].dtype == torch.float32:
                 inputs[key] = inputs[key].to(model.dtype)
@@ -67,41 +44,50 @@ def segment(image: Image.Image, text: str, threshold: float, mask_threshold: flo
         
         n_masks = len(results['masks'])
         if n_masks == 0:
-            return image, f"❌ No objects found matching '{text}' (try adjusting thresholds or changing prompt)."
+            return (image, []), f"❌ No objects found matching '{text}' (try adjusting thresholds)."
         
-        overlaid_image = overlay_masks(image, results["masks"])
+        # Format for AnnotatedImage: list of (mask, label) tuples
+        # mask should be numpy array with values 0-1 (float) matching image dimensions
+        annotations = []
+        for i, (mask, score) in enumerate(zip(results['masks'], results['scores'])):
+            # Convert binary mask to float numpy array (0-1 range)
+            mask_np = mask.cpu().numpy().astype(np.float32)
+            label = f"{text} #{i+1} ({score:.2f})"
+            annotations.append((mask_np, label))
         
-        scores_text = ", ".join([f"{s:.2f}" for s in results['scores'].cpu().numpy()[:5]])  # Top 5 scores
+        scores_text = ", ".join([f"{s:.2f}" for s in results['scores'].cpu().numpy()[:5]])
         info = f"✅ Found **{n_masks}** objects matching **'{text}'**\nConfidence scores: {scores_text}{'...' if n_masks > 5 else ''}"
         
-        return overlaid_image, info
+        # Return tuple: (base_image, list_of_annotations)
+        return (image, annotations), info
         
     except Exception as e:
-        return image, f"❌ Error during segmentation: {str(e)}"
+        return (image, []), f"❌ Error during segmentation: {str(e)}"
 
 def clear_all():
     """Clear all inputs and outputs"""
-    return None, "", None, 0.5, 0.5
+    return None, "", None, 0.5, 0.5, "📝 Enter a prompt and click **Segment** to start."
 
 def segment_example(image_path: str, prompt: str):
     """Handle example clicks"""
-    image = Image.open(image_path) if image_path else None
+    if image_path.startswith("http"):
+        image = Image.open(requests.get(image_path, stream=True).raw).convert("RGB")
+    else:
+        image = Image.open(image_path).convert("RGB")
     return segment(image, prompt, 0.5, 0.5)
 
 # Gradio Interface
 with gr.Blocks(
     theme=gr.themes.Soft(),
     title="SAM3 - Promptable Concept Segmentation",
-    css="""
-    .gradio-container {max-width: 1400px !important;}
-    """
+    css=".gradio-container {max-width: 1400px !important;}"
 ) as demo:
     gr.Markdown(
         """
         # SAM3 - Promptable Concept Segmentation (PCS)
         
-        **SAM3** performs zero-shot instance segmentation using natural language prompts on images.
-        Upload an image, enter a text prompt (e.g., "person", "car", "dog"), and get segmentation masks for all matching objects.
+        **SAM3** performs zero-shot instance segmentation using natural language prompts.
+        Upload an image, enter a text prompt (e.g., "person", "car", "dog"), and get segmentation masks.
         
         Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
         """
@@ -114,15 +100,17 @@ with gr.Blocks(
             type="pil",
             height=400,
         )
+        # AnnotatedImage expects: (base_image, [(mask, label), ...])
         image_output = gr.AnnotatedImage(
             label="Output (Segmented Image)",
             height=400,
+            show_legend=True,
         )
     
     with gr.Row():
         text_input = gr.Textbox(
             label="Text Prompt",
-            placeholder="e.g., a person, ear, cat, bicycle...",
+            placeholder="e.g., person, ear, cat, bicycle...",
             scale=3
         )
         clear_btn = gr.Button("🔍 Clear", size="sm", variant="secondary")
@@ -134,7 +122,7 @@ with gr.Blocks(
             value=0.5,
             step=0.01,
             label="Detection Threshold",
-            info="Higher values = fewer detections (objectness confidence)"
+            info="Higher = fewer detections"
         )
         mask_thresh_slider = gr.Slider(
             minimum=0.0,
@@ -142,7 +130,7 @@ with gr.Blocks(
             value=0.5,
             step=0.01,
             label="Mask Threshold",
-            info="Higher values = sharper masks"
+            info="Higher = sharper masks"
         )
     
     info_output = gr.Markdown(
@@ -152,7 +140,6 @@ with gr.Blocks(
     
     segment_btn = gr.Button("🎯 Segment", variant="primary", size="lg")
     
-    # Add some example prompts
     gr.Examples(
         examples=[
             ["http://images.cocodataset.org/val2017/000000077595.jpg", "cat"],
@@ -163,29 +150,23 @@ with gr.Blocks(
         cache_examples=True,
     )
     
-    # Clear button handler
     clear_btn.click(
         fn=clear_all,
-        outputs=[image_input, text_input, image_output, thresh_slider, mask_thresh_slider]
+        outputs=[image_input, text_input, image_output, thresh_slider, mask_thresh_slider, info_output]
     )
     
-    # Segment button handler
     segment_btn.click(
         fn=segment,
         inputs=[image_input, text_input, thresh_slider, mask_thresh_slider],
         outputs=[image_output, info_output]
-    ).then(
-        fn=lambda: None,
     )
     
-    
     gr.Markdown(
         """
         ### Notes
         - **Model**: [facebook/sam3](https://huggingface.co/facebook/sam3)
-        - Supports natural language prompts like "a red car" or simple nouns.
-        - GPU recommended for faster inference.
-        - Thresholds control detection sensitivity and mask quality.
+        - Click on segments in the output to see labels
+        - GPU recommended for faster inference
         """
     )