use vit_b_lm by default

bioimage-io · Dec 14, 2024 · b785639 · b785639
1 parent ffba37a
commit b785639
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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/",