fix: medsam model

app.py

@@ -1,4 +1,3 @@
-import os
 import numpy as np
 import gradio as gr
 import torch
@@ -6,8 +5,9 @@ import cv2
 from segment_anything import SamPredictor, sam_model_registry
 
 # Global variables
+MODELS = ["./models/sam_vit_b_01ec64.pth", "./models/medsam_vitb.pth"]
 OFFICIAL_CHECKPOINT = "./models/sam_vit_b_01ec64.pth"
-MEDSAM_CHECKPOINT = "./models/medsam_vitb.pth"
+MEDSAM_CHECKPOINT = "./models/medsam_vitb_best.pth"
 MODEL_TYPE = "vit_b"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -22,6 +22,17 @@ MEDSAM.to(device=DEVICE)
 MEDSAM_PREDICTOR = SamPredictor(MEDSAM)
 
 
+def load_model(model_choice: int) -> SamPredictor:
+    """Load model."""
+    print("model_choice", model_choice)
+    if model_choice == 0:
+        return SAM_PREDICTOR
+    elif model_choice == 1:
+        return MEDSAM_PREDICTOR
+    else:
+        raise ValueError("Model choice must be 0 or 1")
+
+
 def draw_contour(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     # draw contour
     contour_image = image.copy()
@@ -53,12 +64,12 @@ def inference(
     return contour_image
 
 
-def extract_object_by_event(
-    image: np.ndarray, evt: gr.SelectData
-):
+def extract_object_by_event(model_choice: int, image: np.ndarray, evt: gr.SelectData):
+    """Extract object by mouse click."""
+    predictor = load_model(model_choice)
     click_h, click_w = evt.index
 
-    return inference(SAM_PREDICTOR, image, click_h, click_w)
+    return inference(predictor, image, click_h, click_w)
 
 
 def get_coords(evt: gr.SelectData):
@@ -76,10 +87,11 @@ with gr.Blocks() as demo:
         )
     with gr.Row():
         # select model
-        model_type = gr.Dropdown(
-            ["oficial_sam", "breast-cancer-sam"],
-            value="oficial_sam",
+        model_choice = gr.Dropdown(
             label="Select Model",
+            choices=[m for m in MODELS],
+            type="index",
+            interactive=True,
         )
 
     # Segment image
@@ -96,7 +108,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         coord_h = gr.Number(label="Mouse coords h")
         coord_w = gr.Number(label="Mouse coords w")
-    input_image.select(extract_object_by_event, [input_image], output)
+    input_image.select(extract_object_by_event, [model_choice, input_image], output)
     input_image.select(get_coords, None, [coord_h, coord_w])
 
 demo.queue().launch(debug=True, enable_queue=True)

models/{medsam_vitb.pth → medsam_vitb_best.pth}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912
-size 375042383
+oid sha256:2b194d15811c1471edd68a65b19d76932f49450cb8e5f20cddae08ac142c8101
+size 375063985