add examples and improve SAM2AutomaticMaskGenerator results

README.md

@@ -1,8 +1,8 @@
 ---
 title: Segment Anything 2
-emoji: 🦄
-colorFrom: pink
-colorTo: red
+emoji: 🔥
+colorFrom: purple
+colorTo: green
 sdk: gradio
 sdk_version: 4.39.0
 app_file: app.py

app.py

@@ -31,6 +31,11 @@ Segment Anything Model 2 (SAM 2) is a foundation model designed to address promp
 visual segmentation in both images and videos. **Video segmentation will be available 
 soon.**
 """
+EXAMPLES = [
+    ["tiny", MASK_GENERATION_MODE, "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", None],
+    ["tiny", MASK_GENERATION_MODE, "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", None],
+    ["tiny", MASK_GENERATION_MODE, "https://media.roboflow.com/notebooks/examples/dog-4.jpeg", None],
+]
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
@@ -102,6 +107,19 @@ with gr.Blocks() as demo:
                 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,
+                mode_dropdown_component,
+                image_input_component,
+                image_prompter_input_component,
+            ],
+            outputs=[image_output_component],
+            run_on_click=True
+        )
 
 
     def on_mode_dropdown_change(text):

utils/models.py

@@ -27,5 +27,14 @@ def load_models(
     for key, (config, checkpoint) in CHECKPOINTS.items():
         model = build_sam2(config, checkpoint, device=device)
         image_predictors[key] = SAM2ImagePredictor(sam_model=model)
-        mask_generators[key] = SAM2AutomaticMaskGenerator(model=model)
+        mask_generators[key] = SAM2AutomaticMaskGenerator(
+            model=model,
+            points_per_side=32,
+            points_per_batch=64,
+            pred_iou_thresh=0.7,
+            stability_score_thresh=0.92,
+            stability_score_offset=0.7,
+            crop_n_layers=1,
+            box_nms_thresh=0.7,
+        )
     return image_predictors, mask_generators