feat: add medsam

app.py

@@ -5,12 +5,21 @@ import torch
 import cv2
 from segment_anything import SamPredictor, sam_model_registry
 
-CHECKPOINT = "./models/sam_vit_h_4b8939.pth"
-MODEL_TYPE = "vit_h"
+# Global variables
+OFFICIAL_CHECKPOINT = "./models/sam_vit_b_01ec64.pth"
+MEDSAM_CHECKPOINT = "./models/medsam_vitb.pth"
+MODEL_TYPE = "vit_b"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-SAM = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT)
+
+# Model
+## OFFICIAL SAM
+SAM = sam_model_registry[MODEL_TYPE](checkpoint=OFFICIAL_CHECKPOINT)
 SAM.to(device=DEVICE)
 SAM_PREDICTOR = SamPredictor(SAM)
+## MEDSAM
+MEDSAM = sam_model_registry[MODEL_TYPE](checkpoint=MEDSAM_CHECKPOINT)
+MEDSAM.to(device=DEVICE)
+MEDSAM_PREDICTOR = SamPredictor(MEDSAM)
 
 
 def draw_contour(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
@@ -23,17 +32,20 @@ def draw_contour(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
     return contour_image, contours
 
 
-def inference(image: np.ndarray, coord_y: int, coord_x: int) -> np.ndarray:
+def inference(
+    predictor: SamPredictor, image: np.ndarray, coord_y: int, coord_x: int
+) -> np.ndarray:
     """Inference."""
-    SAM_PREDICTOR.set_image(image)
+    predictor.set_image(image)
 
     input_point = np.array([[coord_y, coord_x]])
     input_label = np.array([1])
-    mask, _, _ = SAM_PREDICTOR.predict(
+    mask, _, _ = predictor.predict(
         point_coords=input_point,
         point_labels=input_label,
         multimask_output=False,
     )
+
     h, w = mask.shape[-2:]
     mask = mask.reshape(h, w, 1)
     mask = (mask * 255).astype(np.uint8)
@@ -63,7 +75,7 @@ with gr.Blocks() as demo:
         )
 
     # Segment image
-    with gr.Tab(label="Image"):
+    with gr.Tab(label="SAM Inference"):
         with gr.Row().style(equal_height=True):
             with gr.Column(label="Input Image"):
                 # input image
@@ -80,7 +92,12 @@ with gr.Blocks() as demo:
     input_image.select(get_coords, None, [coord_h, coord_w])
     gr.Examples(
         examples=[
-            [os.path.join(os.path.dirname(__file__), "samples/bears.jpg"), 1300, 950]
+            [os.path.join(os.path.dirname(__file__), "samples/bears.jpg"), 1300, 950],
+            [
+                os.path.join(os.path.dirname(__file__), "samples/breast_cancer.png"),
+                125,
+                60,
+            ],
         ],
         inputs=[input_image, coord_h, coord_w],
         outputs=output,

models/{sam_vit_h_4b8939.pth → medsam_vitb.pth}

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

models/sam_vit_b_01ec64.pth

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

requirements.txt

@@ -1,5 +1,6 @@
 opencv-python 
 matplotlib
 gradio
+transformers
 torch
 torchvision

scripts/example.py

@@ -16,6 +16,7 @@ device = "cpu"
 # Load model
 sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
 sam.to(device=device)
+
 predictor = SamPredictor(sam)
 
 # Preprocessing the image