Initial Commit

.gitattributes

@@ -34,3 +34,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 weights/sam_vit_b_01ec64.pth filter=lfs diff=lfs merge=lfs -text
+weights/sam_vit_h_4b8939.pth filter=lfs diff=lfs merge=lfs -text
+weights/sam_vit_l_0b3195.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -6,17 +6,26 @@ from segment_anything import sam_model_registry, SamPredictor
 from preprocess import show_mask, show_points, show_box
 import gradio as gr
 
-sam_checkpoint = "weights/sam_vit_b_01ec64.pth"
-model_type = "vit_b"
+sam_checkpoint = {
+                    "ViT-base": "weights/sam_vit_b_01ec64.pth",
+                    "ViT-large": "weights/sam_vit_l_0b3195.pth",
+                    "ViT-huge": "weights/sam_vit_h_4b8939.pth",
+                 }
+model_type = {
+                "ViT-base": "vit_b",
+                "ViT-large": "vit_l",
+                "ViT-huge": "vit_h",
+             }
 device = "cuda" if torch.cuda.is_available() else "cpu"
-sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
-sam.to(device=device)
-predictor = SamPredictor(sam)
+
 
 def get_coords(evt: gr.SelectData):
     return f"{evt.index[0]}, {evt.index[1]}"
 
-def inference(image, input_label):
+def inference(image, input_label, model_choice):
+    sam = sam_model_registry[model_type[model_choice]](checkpoint=sam_checkpoint[model_choice])
+    sam.to(device=device)
+    predictor = SamPredictor(sam)
     predictor.set_image(image)
     input_point = np.array([[int(input_label['label'].split(',')[0]), int(input_label['label'].split(',')[1])]])
     input_label = np.array([1])
@@ -38,8 +47,10 @@ with my_app:
         with gr.TabItem("Select your image"):
             with gr.Row():
                 with gr.Column():
-                    img_source = gr.Image(label="Please select picture.", value='./images/truck.jpg', shape=(1024, 1024))
+                    img_source = gr.Image(label="Please select picture and click the part to segment", 
+                                          value='./images/truck.jpg', shape=(1024, 1024))
                     coords = gr.Label(label="Image Coordinate")
+                    model_choice = gr.Dropdown(['ViT-base', 'ViT-large', 'ViT-huge'], label='Model Backbone')
                     infer = gr.Button(label="Segment")
                 with gr.Column():
                     img_output = gr.Image(label="Output Mask", shape=(1024, 1024))
@@ -49,7 +60,8 @@ with my_app:
             inference,
             [
                 img_source,
-                coords
+                coords,
+                model_choice
             ],
             [
                 img_output

weights/sam_vit_h_4b8939.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e
+size 2564550879

weights/sam_vit_l_0b3195.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622
+size 1249524607