Enable passing custom predictor to all annotator functions

Author: constantinpape

micro_sam/sam_annotator/annotator_3d.py

@@ -7,6 +7,7 @@
 from magicgui import magicgui
 from napari import Viewer
 from napari.utils import progress
+from segment_anything import SamPredictor
 
 from .. import util
 from ..prompt_based_segmentation import segment_from_mask
@@ -195,10 +196,15 @@ def annotator_3d(
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
     return_viewer: bool = False,
+    predictor: Optional[SamPredictor] = None,
 ) -> None:
     # for access to the predictor and the image embeddings in the widgets
     global PREDICTOR, IMAGE_EMBEDDINGS
-    PREDICTOR = util.get_sam_model(model_type=model_type)
+
+    if predictor is None:
+        PREDICTOR = util.get_sam_model(model_type=model_type)
+    else:
+        PREDICTOR = predictor
     IMAGE_EMBEDDINGS = util.precompute_image_embeddings(
         PREDICTOR, raw, save_path=embedding_path, tile_shape=tile_shape, halo=halo,
         wrong_file_callback=show_wrong_file_warning,

micro_sam/sam_annotator/annotator_tracking.py

@@ -9,6 +9,7 @@
 from napari import Viewer
 from napari.utils import progress
 from scipy.ndimage import shift
+from segment_anything import SamPredictor
 
 # this is more precise for comuting the centers, but slow!
 # from vigra.filters import eccentricityCenters
@@ -355,12 +356,16 @@ def annotator_tracking(
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
     return_viewer: bool = False,
+    predictor: Optional[SamPredictor] = None,
 ) -> None:
     # global state
     global PREDICTOR, IMAGE_EMBEDDINGS, CURRENT_TRACK_ID, LINEAGE
     global TRACKING_WIDGET
 
-    PREDICTOR = util.get_sam_model(model_type=model_type)
+    if predictor is None:
+        PREDICTOR = util.get_sam_model(model_type=model_type)
+    else:
+        PREDICTOR = predictor
     IMAGE_EMBEDDINGS = util.precompute_image_embeddings(
         PREDICTOR, raw, save_path=embedding_path, tile_shape=tile_shape, halo=halo,
         wrong_file_callback=show_wrong_file_warning,

micro_sam/sam_annotator/image_series_annotator.py

@@ -8,6 +8,8 @@
 
 from magicgui import magicgui
 from napari.utils import progress as tqdm
+from segment_anything import SamPredictor
+
 from .annotator_2d import annotator_2d
 from .. import util
 
@@ -32,6 +34,7 @@ def image_series_annotator(
     image_files: List[str],
     output_folder: str,
     embedding_path: Optional[str] = None,
+    predictor: Optional[SamPredictor] = None,
     **kwargs
 ) -> None:
     """
@@ -45,7 +48,8 @@ def image_series_annotator(
     os.makedirs(output_folder, exist_ok=True)
     next_image_id = 0
 
-    predictor = util.get_sam_model(model_type=kwargs.get("model_type", "vit_h"))
+    if predictor is None:
+        predictor = util.get_sam_model(model_type=kwargs.get("model_type", "vit_h"))
     if embedding_path is None:
         embedding_paths = None
     else:
@@ -101,12 +105,13 @@ def image_folder_annotator(
     output_folder: str,
     pattern: str = "*",
     embedding_path: Optional[str] = None,
+    predictor: Optional[SamPredictor] = None,
     **kwargs
 ) -> None:
     """
     """
     image_files = sorted(glob(os.path.join(input_folder, pattern)))
-    image_series_annotator(image_files, output_folder, embedding_path, **kwargs)
+    image_series_annotator(image_files, output_folder, embedding_path, predictor, **kwargs)
 
 
 def main():