initial version of SAM2 space

app.py

@@ -6,27 +6,48 @@ import supervision as sv
 import torch
 from PIL import Image
 from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
-from sam2.build_sam import build_sam2
+
+from utils.models import load_models, CHECKPOINT_NAMES
 
 MARKDOWN = """
 # Segment Anything Model 2 🔥
+<div>
+    <a href="https://github.com/facebookresearch/segment-anything-2">
+        <img src="https://badges.aleen42.com/src/github.svg" alt="GitHub" style="display:inline-block;">
+    </a>
+    <a href="https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-segment-images-with-sam-2.ipynb">
+        <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Colab" style="display:inline-block;">
+    </a>
+    <a href="https://blog.roboflow.com/what-is-segment-anything-2/">
+        <img src="https://raw.githubusercontent.com/roboflow-ai/notebooks/main/assets/badges/roboflow-blogpost.svg" alt="Roboflow" style="display:inline-block;">
+    </a>
+    <a href="https://www.youtube.com/watch?v=Dv003fTyO-Y">
+        <img src="https://badges.aleen42.com/src/youtube.svg" alt="YouTube" style="display:inline-block;">
+    </a>
+</div>
 
 Segment Anything Model 2 (SAM 2) is a foundation model designed to address promptable 
 visual segmentation in both images and videos. The model extends its functionality to 
 video by treating images as single-frame videos. Its design, a simple transformer 
-architecture with streaming memory, enables real-time video processing.
+architecture with streaming memory, enables real-time video processing. A 
+model-in-the-loop data engine, which enhances the model and data through user 
+interaction, was built to collect the SA-V dataset, the largest video segmentation 
+dataset to date. SAM 2, trained on this extensive dataset, delivers robust performance 
+across diverse tasks and visual domains.
 """
+EXAMPLES = [
+    ["tiny", "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"],
+    ["small", "https://media.roboflow.com/notebooks/examples/dog-3.jpeg"],
+    ["large", "https://media.roboflow.com/notebooks/examples/dog-3.jpeg"],
+]
 
-# DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-DEVICE = torch.device('cuda')
-
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
-CHECKPOINT = "checkpoints/sam2_hiera_large.pt"
-CONFIG = "sam2_hiera_l.yaml"
+MODELS = load_models(device=DEVICE)
 
-sam2_model = build_sam2(CONFIG, CHECKPOINT, device=DEVICE, apply_postprocessing=False)
 
-def process(image_input) -> Optional[Image.Image]:
+def process(checkpoint_dropdown, image_input) -> Optional[Image.Image]:
+    sam2_model = MODELS[checkpoint_dropdown]
     mask_generator = SAM2AutomaticMaskGenerator(sam2_model)
     image = np.array(image_input.convert("RGB"))
     sam_result = mask_generator.generate(image)
@@ -36,17 +57,31 @@ def process(image_input) -> Optional[Image.Image]:
 
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN)
-
+    with gr.Row():
+        checkpoint_dropdown_component = gr.Dropdown(
+            choices=CHECKPOINT_NAMES,
+            value=CHECKPOINT_NAMES[0],
+            label="Checkpoint", info="Select a SAM2 checkpoint to use.",
+            interactive=True
+        )
     with gr.Row():
         with gr.Column():
             image_input_component = gr.Image(type='pil', label='Upload image')
             submit_button_component = gr.Button(value='Submit', variant='primary')
         with gr.Column():
             image_output_component = gr.Image(type='pil', label='Image Output')
+    with gr.Row():
+        gr.Examples(
+            fn=process,
+            examples=EXAMPLES,
+            inputs=[checkpoint_dropdown_component, image_input_component],
+            outputs=[image_output_component],
+            run_on_click=True
+        )
 
     submit_button_component.click(
         fn=process,
-        inputs=[image_input_component],
+        inputs=[checkpoint_dropdown_component, image_input_component],
         outputs=[image_output_component]
     )
 

utils/models.py

@@ -1,6 +1,20 @@
+import torch
+
+from typing import Dict, Any
+from sam2.build_sam import build_sam2
+
+CHECKPOINT_NAMES = ["tiny", "small", "base_plus", "large"]
+
 CHECKPOINTS = {
     "tiny": ["sam2_hiera_t.yaml", "checkpoints/sam2_hiera_tiny.pt"],
     "small": ["sam2_hiera_s.yaml", "checkpoints/sam2_hiera_small.pt"],
     "base_plus": ["sam2_hiera_b+.yaml", "checkpoints/sam2_hiera_base_plus.pt"],
     "large": ["sam2_hiera_l.yaml", "checkpoints/sam2_hiera_large.pt"],
-}
+}
+
+
+def load_models(device: torch.device) -> Dict[str, Any]:
+    models = {}
+    for key, (config, checkpoint) in CHECKPOINTS.items():
+        models[key] = build_sam2(config, checkpoint, device=device, apply_postprocessing=False)
+    return models