add segmentation function

app.py

@@ -1,9 +1,65 @@
+import os
+import urllib
+
+import cv2
 import gradio as gr
+import matplotlib.pyplot as plt
+import numpy as np
+
+from PIL import Image
+
+from segment_anything import SamAutomaticMaskGenerator, sam_model_registry
+
+
+# download model weights
+ckpts_dir = os.path.join(os.getcwd() + "/ckpts")
+if not os.path.exists(ckpts_dir):
+    os.makedirs(ckpts_dir)
+ckpt_path = os.path.join(ckpts_dir + "/vit_b.pth")
+if not os.path.exists(ckpt_path):
+    url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+    urllib.request.urlretrieve(url, filename=ckpt_path)
 
+# setup model
+sam = sam_model_registry["vit_b"](checkpoint=ckpt_path)
+mask_generator = SamAutomaticMaskGenerator(sam)
 
+# copied from: https://github.com/facebookresearch/segment-anything
+def show_anns(anns):
+    if len(anns) == 0:
+        return
+    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
+    ax = plt.gca()
+    ax.set_autoscale_on(False)
+    polygons = []
+    color = []
+    for ann in sorted_anns:
+        m = ann['segmentation']
+        img = np.ones((m.shape[0], m.shape[1], 3))
+        color_mask = np.random.random((1, 3)).tolist()[0]
+        for i in range(3):
+            img[:,:,i] = color_mask[i]
+        ax.imshow(np.dstack((img, m*0.35)))
+
+# demo function
 def segment_image(input_image):
-    output_image = input_image
-    return output_image
+
+    # generate masks
+    masks = mask_generator.generate(input_image)
+
+    # add masks to image
+    plt.clf()
+    ppi = 100
+    height, width, _ = input_image.shape
+    plt.figure(figsize=(width / ppi, height / ppi))  # convert pixel to inches
+    plt.imshow(input_image)
+    show_anns(masks)
+    plt.axis('off')
+
+    # save and get figure
+    plt.savefig('output_figure.png', bbox_inches='tight')
+    output_image = cv2.imread('output_figure.png')
+    return Image.fromarray(output_image)
 
 
 with gr.Blocks() as demo: