no shrooms

sam2_mask.py

@@ -12,10 +12,9 @@ MODEL = "facebook/sam2.1-hiera-large"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 @spaces.GPU()
+
 def predict_masks(image, points):
     """Predict a single mask from the image based on selected points."""
-    global PREDICTOR
-    PREDICTOR = SAM2ImagePredictor.from_pretrained(MODEL, device=DEVICE)
     image_np = np.array(image)
     points_list = [[point["x"], point["y"]] for point in points]
     input_labels = [1] * len(points_list)
@@ -37,8 +36,15 @@ def predict_masks(image, points):
     else:
         return image_np
 
+
 def create_sam2_mask_interface():
     """Create the Gradio interface for SAM2 mask generation."""
+    def update_mask(prompts):
+        """Update the mask based on the prompts."""
+        image = prompts["image"]
+        points = prompts["points"]
+        return predict_masks(image, points)
+
     with gr.Blocks() as sam2_mask_tab:
         gr.Markdown("# Object Segmentation with SAM2")
         gr.Markdown(
@@ -49,15 +55,14 @@ def create_sam2_mask_interface():
         with gr.Row():
             with gr.Column():
                 upload_image_input = ImagePrompter(show_label=False)
-                submit_button = gr.Button("Submit")
             with gr.Column():
                 image_output = gr.Image(label="Segmented Image", type="pil", height=400)
         
-        # Define the action triggered by the submit button
-        submit_button.click(
-            fn=predict_masks,
-            inputs=[upload_image_input.image, upload_image_input.points],
-            outputs=image_output,
+        # Define the action triggered by the upload_image_input change
+        upload_image_input.change(
+            fn=update_mask,
+            inputs=[upload_image_input],
+            outputs=[image_output],
             show_progress=True,
         )