points per side argument added

README.md

@@ -1,7 +1,7 @@
 ---
-title: Segment Anything Model 2
-emoji: 💻
-colorFrom: blue
+title: Segment Anything 2
+emoji: 🦄
+colorFrom: pink
 colorTo: red
 sdk: gradio
 sdk_version: 4.39.0

app.py

@@ -36,9 +36,10 @@ dataset to date. SAM 2, trained on this extensive dataset, delivers robust perfo
 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"],
+    ["tiny", "https://media.roboflow.com/notebooks/examples/dog-2.jpeg", 16],
+    ["small", "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 16],
+    ["large", "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 16],
+    ["large", "https://media.roboflow.com/notebooks/examples/dog-3.jpeg", 64],
 ]
 
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -46,9 +47,11 @@ MASK_ANNOTATOR = sv.MaskAnnotator(color_lookup=sv.ColorLookup.INDEX)
 MODELS = load_models(device=DEVICE)
 
 
-def process(checkpoint_dropdown, image_input) -> Optional[Image.Image]:
-    sam2_model = MODELS[checkpoint_dropdown]
-    mask_generator = SAM2AutomaticMaskGenerator(sam2_model)
+def process(checkpoint_dropdown, image_input, points_per_side) -> Optional[Image.Image]:
+    model = MODELS[checkpoint_dropdown]
+    mask_generator = SAM2AutomaticMaskGenerator(
+        model=model,
+        points_per_side=points_per_side)
     image = np.array(image_input.convert("RGB"))
     sam_result = mask_generator.generate(image)
     detections = sv.Detections.from_sam(sam_result=sam_result)
@@ -64,6 +67,14 @@ with gr.Blocks() as demo:
             label="Checkpoint", info="Select a SAM2 checkpoint to use.",
             interactive=True
         )
+        points_per_side_component = gr.Slider(
+            minimum=16,
+            maximum=64,
+            value=16,
+            step=16,
+            label="Points per side",
+            info="the number of points to be sampled along one side of the image."
+        )
     with gr.Row():
         with gr.Column():
             image_input_component = gr.Image(type='pil', label='Upload image')
@@ -74,14 +85,22 @@ with gr.Blocks() as demo:
         gr.Examples(
             fn=process,
             examples=EXAMPLES,
-            inputs=[checkpoint_dropdown_component, image_input_component],
+            inputs=[
+                checkpoint_dropdown_component,
+                image_input_component,
+                points_per_side_component
+            ],
             outputs=[image_output_component],
             run_on_click=True
         )
 
     submit_button_component.click(
         fn=process,
-        inputs=[checkpoint_dropdown_component, image_input_component],
+        inputs=[
+            checkpoint_dropdown_component,
+            image_input_component,
+            points_per_side_component
+        ],
         outputs=[image_output_component]
     )