Add option to use fine-tuned models in example scripts

Author: constantinpape

examples/sam_annotator_2d.py

@@ -3,47 +3,72 @@
 from micro_sam.sample_data import fetch_hela_2d_example_data, fetch_livecell_example_data, fetch_wholeslide_example_data
 
 
-def livecell_annotator():
+def livecell_annotator(use_finetuned_model):
     """Run the 2d annotator for an example image from the LiveCELL dataset.
 
     See https://doi.org/10.1038/s41592-021-01249-6 for details on the data.
     """
     example_data = fetch_livecell_example_data("./data")
     image = imageio.imread(example_data)
-    embedding_path = "./embeddings/embeddings-livecell.zarr"
-    annotator_2d(image, embedding_path, show_embeddings=False)
 
+    if use_finetuned_model:
+        embedding_path = "./embeddings/embeddings-livecell-vit_h_lm.zarr"
+        model_type = "vit_h_lm"
+    else:
+        embedding_path = "./embeddings/embeddings-livecell.zarr"
+        model_type = "vit_h"
 
-def hela_2d_annotator():
+    annotator_2d(image, embedding_path, show_embeddings=False, model_type=model_type)
+
+
+def hela_2d_annotator(use_finetuned_model):
     """Run the 2d annotator for an example image form the cell tracking challenge HeLa 2d dataset.
     """
     example_data = fetch_hela_2d_example_data("./data")
     image = imageio.imread(example_data)
-    embedding_path = "./embeddings/embeddings-hela2d.zarr"
-    annotator_2d(image, embedding_path, show_embeddings=False)
+
+    if use_finetuned_model:
+        embedding_path = "./embeddings/embeddings-hela2d-vit_h_lm.zarr"
+        model_type = "vit_h_lm"
+    else:
+        embedding_path = "./embeddings/embeddings-hela2d.zarr"
+        model_type = "vit_h"
+
+    annotator_2d(image, embedding_path, show_embeddings=False, model_type=model_type)
 
 
-def wholeslide_annotator():
+def wholeslide_annotator(use_finetuned_model):
     """Run the 2d annotator with tiling for an example whole-slide image from the
     NeuRIPS cell segmentation challenge.
 
     See https://neurips22-cellseg.grand-challenge.org/ for details on the data.
     """
     example_data = fetch_wholeslide_example_data("./data")
     image = imageio.imread(example_data)
-    embedding_path = "./embeddings/whole-slide-embeddings.zarr"
-    annotator_2d(image, embedding_path, tile_shape=(1024, 1024), halo=(256, 256))
+
+    if use_finetuned_model:
+        embedding_path = "./embeddings/whole-slide-embeddings-vit_h_lm.zarr"
+        model_type = "vit_h_lm"
+    else:
+        embedding_path = "./embeddings/whole-slide-embeddings.zarr"
+        model_type = "vit_h"
+
+    annotator_2d(image, embedding_path, tile_shape=(1024, 1024), halo=(256, 256), model_type=model_type)
 
 
 def main():
+    # whether to use the fine-tuned SAM model
+    # this feature is still experimental!
+    use_finetuned_model = False
+
     # 2d annotator for livecell data
-    # livecell_annotator()
+    # livecell_annotator(use_finetuned_model)
 
     # 2d annotator for cell tracking challenge hela data
-    hela_2d_annotator()
+    # hela_2d_annotator(use_finetuned_model)
 
     # 2d annotator for a whole slide image
-    # wholeslide_annotator()
+    wholeslide_annotator(use_finetuned_model)
 
 
 if __name__ == "__main__":

examples/sam_annotator_tracking.py

@@ -3,20 +3,31 @@
 from micro_sam.sample_data import fetch_tracking_example_data
 
 
-def track_ctc_data():
+def track_ctc_data(use_finetuned_model):
     """Run interactive tracking for data from the cell tracking challenge.
     """
     # download the example data
     example_data = fetch_tracking_example_data("./data")
     # load the example data (load the sequence of tif files as timeseries)
     with open_file(example_data, mode="r") as f:
         timeseries = f["*.tif"]
+
+    if use_finetuned_model:
+        embedding_path = "./embeddings/embeddings-ctc-vit_h_lm.zarr"
+        model_type = "vit_h_lm"
+    else:
+        embedding_path = "./embeddings/embeddings-ctc.zarr"
+        model_type = "vit_h"
+
     # start the annotator with cached embeddings
-    annotator_tracking(timeseries, embedding_path="./embeddings/embeddings-ctc.zarr", show_embeddings=False)
+    annotator_tracking(timeseries, embedding_path=embedding_path, show_embeddings=False, model_type=model_type)
 
 
 def main():
-    track_ctc_data()
+    # whether to use the fine-tuned SAM model
+    # this feature is still experimental!
+    use_finetuned_model = False
+    track_ctc_data(use_finetuned_model)
 
 
 if __name__ == "__main__":