demo file with dependency

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+import numpy as np
+import torch
+from torchvision.transforms import ToTensor
+from PIL import Image
+
+# loading EfficientSAM model
+model_path = "efficientsam_s_cpu.jit"
+with open(model_path, "rb") as f:
+    model = torch.jit.load(f)
+
+# getting mask using points
+def get_sam_mask_using_points(img_tensor, pts_sampled, model):
+    pts_sampled = torch.reshape(torch.tensor(pts_sampled), [1, 1, -1, 2])
+    max_num_pts = pts_sampled.shape[2]
+    pts_labels = torch.ones(1, 1, max_num_pts)
+
+    predicted_logits, predicted_iou = model(
+        img_tensor[None, ...],
+        pts_sampled,
+        pts_labels,
+    )
+    predicted_logits = predicted_logits.cpu()
+    all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
+    predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
+
+    max_predicted_iou = -1
+    selected_mask_using_predicted_iou = None
+    for m in range(all_masks.shape[0]):
+        curr_predicted_iou = predicted_iou[m]
+        if (
+            curr_predicted_iou > max_predicted_iou
+            or selected_mask_using_predicted_iou is None
+        ):
+            max_predicted_iou = curr_predicted_iou
+            selected_mask_using_predicted_iou = all_masks[m]
+    return selected_mask_using_predicted_iou
+
+# examples
+examples = [["examples/image1.jpg"], ["examples/image2.jpg"], ["examples/image3.jpg"], ["examples/image4.jpg"],
+            ["examples/image5.jpg"], ["examples/image6.jpg"], ["examples/image7.jpg"], ["examples/image8.jpg"],
+            ["examples/image9.jpg"], ["examples/image10.jpg"], ["examples/image11.jpg"], ["examples/image12.jpg"]
+            ["examples/image13.jpg"], ["examples/image14.jpg"]]
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        input_img = gr.Image(label="Input",height=512)
+        output_img = gr.Image(label="Selected Segment",height=512)
+
+    with gr.Row():
+        gr.Markdown("Try some of the examples below ⬇️")
+        gr.Examples(examples=examples,
+                    inputs=[input_img])
+
+    def get_select_coords(img, evt: gr.SelectData):
+        img_tensor = ToTensor()(img)
+        _, H, W = img_tensor.shape
+
+        visited_pixels = set()
+        pixels_in_queue = set()
+        pixels_in_segment = set()
+
+        mask = get_sam_mask_using_points(img_tensor, [[evt.index[0], evt.index[1]]], model)
+
+        out = img.copy()
+
+        out = out.astype(np.uint8)
+        out *= mask[:,:,None]
+        for pixel in pixels_in_segment:
+            out[pixel[0], pixel[1]] = img[pixel[0], pixel[1]]
+        print(out)
+        return out
+
+    input_img.select(get_select_coords, [input_img], output_img)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+torch==2.0.1
+torchvision==0.15.2
+gradio
+transformers==4.32.0
+opencv-python
+pandas==2.0.3
+matplotlib==3.7.2