requirements.txt

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image
+import cv2
+from segment_anything import sam_model_registry, SamPredictor
+
+# Load model
+checkpoint = "sam_vit_h_4b8939.pth"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model_type = "vit_h"
+
+sam = sam_model_registry[model_type](checkpoint=checkpoint)
+sam.to(device)
+predictor = SamPredictor(sam)
+
+def segment_image(input_img):
+    np_img = np.array(input_img)
+    image = cv2.cvtColor(np_img, cv2.COLOR_RGB2BGR)
+
+    predictor.set_image(image)
+    h, w, _ = image.shape
+    input_point = np.array([[w // 2, h // 2]])
+    input_label = np.array([1])
+    
+    masks, scores, logits = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=False
+    )
+
+    mask = masks[0].astype(np.uint8) * 255
+    return Image.fromarray(mask)
+
+# UI
+iface = gr.Interface(fn=segment_image,
+                     inputs=gr.Image(type="pil"),
+                     outputs=gr.Image(type="pil"),
+                     title="Segment Anything Model",
+                     description="Upload an image and get a segmentation mask.")
+
+iface.launch()