From b7856397cb2508d9c22a26475011caa383dc3c78 Mon Sep 17 00:00:00 2001
From: Nils Mechtel <nils.mech@gmail.com>
Date: Sat, 14 Dec 2024 23:30:02 +0100
Subject: [PATCH] use vit_b_lm by default

---
 bioimageio_colab/register_sam_service.py      | 4 ++--
 bioimageio_colab/sam.py                       | 2 +-
 plugins/bioimageio-colab-annotator.imjoy.html | 2 +-
 test/test_model_service.py                    | 4 ++--
 test/test_sam.py                              | 2 +-
 5 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/bioimageio_colab/register_sam_service.py b/bioimageio_colab/register_sam_service.py
index 8b12c8d..9fe0845 100644
--- a/bioimageio_colab/register_sam_service.py
+++ b/bioimageio_colab/register_sam_service.py
@@ -251,7 +251,7 @@ def test_model_function(**kwargs: dict):
 
 
 if __name__ == "__main__":
-    model_name = "vit_b"
+    model_name = "vit_b_lm"
     cache_dir = "./model_cache"
     embedding = compute_image_embedding(
         cache_dir=cache_dir,
@@ -261,7 +261,7 @@ def test_model_function(**kwargs: dict):
     )
     mask = compute_mask(
         cache_dir=cache_dir,
-        model_name="vit_b",
+        model_name="vit_b_lm",
         embedding=embedding,
         image_size=(1024, 1024),
         point_coords=np.array([[10, 10]]),
diff --git a/bioimageio_colab/sam.py b/bioimageio_colab/sam.py
index 03644ea..5ed2e23 100644
--- a/bioimageio_colab/sam.py
+++ b/bioimageio_colab/sam.py
@@ -99,6 +99,6 @@ def segment_image(
 
 
 if __name__ == "__main__":
-    sam_predictor = load_model_from_ckpt("vit_b", "./model_cache")
+    sam_predictor = load_model_from_ckpt("vit_b_lm", "./model_cache")
     sam_predictor = compute_embedding(sam_predictor, np.random.rand(1024, 1024))
     masks = segment_image(sam_predictor, [[10, 10]], [1])
diff --git a/plugins/bioimageio-colab-annotator.imjoy.html b/plugins/bioimageio-colab-annotator.imjoy.html
index c3978cf..281fa6c 100644
--- a/plugins/bioimageio-colab-annotator.imjoy.html
+++ b/plugins/bioimageio-colab-annotator.imjoy.html
@@ -29,7 +29,7 @@
         this.imageLayer = null; // Layer displaying the image
         this.annotationLayer = null; // Layer displaying the annotations
         this.edgeColor = "magenta"; // Default edge color for annotations
-        this.modelName = "vit_b"; // Model name for the embeddings
+        this.modelName = "vit_b_lm"; // Model name for the embeddings
     }
 
     async setup() {
diff --git a/test/test_model_service.py b/test/test_model_service.py
index bd548e7..5685eb4 100644
--- a/test/test_model_service.py
+++ b/test/test_model_service.py
@@ -6,7 +6,7 @@
 SERVER_URL = "https://hypha.aicell.io"
 WORKSPACE_NAME = "bioimageio-colab"
 SERVICE_ID = "microsam"
-MODEL_NAME = "vit_b"
+MODEL_NAME = "vit_b_lm"
 
 
 def test_service_functions():
@@ -77,5 +77,5 @@ def test_service_python_api():
     # assert len(polygon_features) == 1  # Only one point given
 
     # Test service test run
-    result = segment_svc.test_model(model_name="vit_b")
+    result = segment_svc.test_model(model_name="vit_b_lm")
     assert result == {"status": "ok"}
diff --git a/test/test_sam.py b/test/test_sam.py
index e8048c8..a83c70b 100644
--- a/test/test_sam.py
+++ b/test/test_sam.py
@@ -7,7 +7,7 @@
 
 
 def test_sam():
-    sam_predictor = load_model_from_ckpt(model_name="vit_b", cache_dir="./model_cache/")
+    sam_predictor = load_model_from_ckpt(model_name="vit_b_lm", cache_dir="./model_cache/")
     assert os.path.exists("./model_cache/sam_vit_b_01ec64.pth")
     image, _ = get_random_image(
         image_folder="./bioimageio_colab/",