Create app.py

app.py

@@ -0,0 +1,52 @@
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+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)
+
+def segment_image(input_image, x, y):
+    # Convert gradio image to PIL Image
+    input_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    
+    # Prepare the image for the model
+    predictor.set_image(input_image)
+    
+    # Prepare the prompt (point)
+    input_point = np.array([[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)
+
+    # Convert the mask to an image
+    mask = masks[0].cpu().numpy()
+    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)
+    
+    return result
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=segment_image,
+    inputs=[
+        gr.Image(type="numpy"),
+        gr.Slider(0, 1000, label="X coordinate"),
+        gr.Slider(0, 1000, label="Y coordinate")
+    ],
+    outputs=gr.Image(type="pil"),
+    title="SAM2 Image Segmentation",
+    description="Upload an image and select a point to segment. Adjust X and Y coordinates to refine the selection."
+)
+
+# Launch the app
+iface.launch()