Create app.py

app.py

@@ -0,0 +1,36 @@
+import torch
+from PIL import Image
+import gradio as gr
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+
+# Load the SAM2 model
+predictor = SAM2ImagePredictor.from_pretrained("facebook/sam2.1-hiera-large")
+
+# Function to predict masks from the image and prompts
+def generate_mask(image, prompt):
+    with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+        predictor.set_image(image)
+        masks, _, _ = predictor.predict(prompt)
+    return masks[0]  # Returning the first mask for simplicity
+
+# Set up the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Image Segmentation using SAM2")
+
+    # Input: Upload an image
+    image_input = gr.Image(label="Upload Image", type="pil")
+
+    # Input: Text prompt for image segmentation
+    prompt_input = gr.Textbox(label="Enter segmentation prompt", placeholder="Describe what you want to segment")
+
+    # Output: Display the mask generated by the SAM2 model
+    output_mask = gr.Image(label="Generated Mask")
+
+    # Button to trigger mask generation
+    generate_button = gr.Button("Generate Mask")
+
+    # Link button click with the segmentation function
+    generate_button.click(fn=generate_mask, inputs=[image_input, prompt_input], outputs=output_mask)
+
+# Launch the Gradio app
+demo.launch()