Create app.py

app.py

@@ -0,0 +1,53 @@
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+from ultralytics import SAM
+
+model = SAM('sam_b.pt')
+
+
+def process_image(input_image):
+    if input_image is None:
+        return None, "Can't find photo", 0
+
+    input_image = np.array(input_image)
+    bgr_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
+
+    r = model(bgr_image)[0]
+    masks = r.masks.xy
+
+    # point = (x, y)
+    # segment_index = point_in_mask(masks, point)
+
+    im_array = r.plot().astype(np.uint8)
+    im_array = cv2.cvtColor(im_array, cv2.COLOR_BGR2RGB)
+    im_pil = Image.fromarray(im_array)
+
+    # cv2.circle(im_array, point, radius=5, color=(255, 255, 0), thickness=2)
+
+    return im_pil
+
+
+def point_in_mask(masks, point):
+    x, y = point
+    for i in range(len(masks) - 1, -1, -1):
+        mask = np.array(masks[i])
+        min_x = min(mask[:, 0])
+        max_x = max(mask[:, 0])
+        min_y = min(mask[:, 1])
+        max_y = max(mask[:, 1])
+        if min_x <= x <= max_x and min_y <= y <= max_y:
+            return i
+    return None
+
+
+iface = gr.Interface(
+    fn=process_image,
+    inputs=[gr.Image()],
+    outputs=[gr.Image(type="numpy")],
+    title="Image Segmenter",
+    description="Upload an image to segment."
+)
+
+iface.launch()