Update app.py

app.py

@@ -2,18 +2,20 @@ import gradio as gr
 import torch
 from PIL import Image
 import numpy as np
-from sam2.sam2_image_predictor import SAM2ImagePredictor
+from sam2 import build_sam2, SamPredictor
 from huggingface_hub import hf_hub_download
 
 # Download the model weights
 model_path = hf_hub_download(repo_id="facebook/sam2-hiera-large", filename="sam2_hiera_large.pth")
 
-# Initialize the SAM2 predictor
-predictor = SAM2ImagePredictor.from_pretrained(model_path)
+# Initialize the SAM2 model
+device = "cpu"  # Use CPU
+model = build_sam2(checkpoint=model_path).to(device)
+predictor = SamPredictor(model)
 
 def segment_image(input_image, x, y):
-    # Convert gradio image to PIL Image
-    input_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    # Convert gradio image to numpy array
+    input_image = np.array(input_image)
     
     # Prepare the image for the model
     predictor.set_image(input_image)
@@ -23,15 +25,18 @@ def segment_image(input_image, x, y):
     input_label = np.array([1])  # 1 for foreground
 
     # Generate the mask
-    with torch.inference_mode():
-        masks, _, _ = predictor.predict(point_coords=input_point, point_labels=input_label)
+    masks, _, _ = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=False,
+    )
 
     # Convert the mask to an image
-    mask = masks[0].cpu().numpy()
+    mask = masks[0]
     mask_image = Image.fromarray((mask * 255).astype(np.uint8))
     
     # Apply the mask to the original image
-    result = Image.composite(input_image, Image.new('RGB', input_image.size, 'black'), mask_image)
+    result = Image.composite(Image.fromarray(input_image), Image.new('RGB', mask_image.size, 'black'), mask_image)
     
     return result
 
@@ -39,7 +44,7 @@ def segment_image(input_image, x, y):
 iface = gr.Interface(
     fn=segment_image,
     inputs=[
-        gr.Image(type="numpy"),
+        gr.Image(type="pil"),
         gr.Slider(0, 1000, label="X coordinate"),
         gr.Slider(0, 1000, label="Y coordinate")
     ],