Merge pull request #98 from computational-cell-analytics/custom_models

Add preliminary custom models

Author: constantinpape

environment_cpu.yaml

@@ -9,7 +9,8 @@ dependencies:
     - pooch
     - python-elf >=0.4.8
     - pytorch
+    - segment-anything
     - torchvision
     - tqdm
-    - pip:
-        - git+https://github.com/facebookresearch/segment-anything.git
+    # - pip:
+    #     - git+https://github.com/facebookresearch/segment-anything.git

environment_gpu.yaml

@@ -10,7 +10,8 @@ dependencies:
     - python-elf >=0.4.8
     - pytorch
     - pytorch-cuda>=11.7  # you may need to update the cuda version to match your system
+    - segment-anything
     - torchvision
     - tqdm
-    - pip:
-        - git+https://github.com/facebookresearch/segment-anything.git
+    # - pip:
+    #     - git+https://github.com/facebookresearch/segment-anything.git

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__":

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():

micro_sam/util.py

@@ -25,13 +25,19 @@
 _MODEL_URLS = {
     "vit_h": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
     "vit_l": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
-    "vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
+    "vit_b": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",
+    # preliminary finetuned models
+    "vit_h_lm": "https://owncloud.gwdg.de/index.php/s/CnxBvsdGPN0TD3A/download",
+    "vit_b_lm": "https://owncloud.gwdg.de/index.php/s/gGlR1LFsav0eQ2k/download",
 }
 _CHECKPOINT_FOLDER = os.environ.get("SAM_MODELS", os.path.expanduser("~/.sam_models"))
 _CHECKSUMS = {
     "vit_h": "a7bf3b02f3ebf1267aba913ff637d9a2d5c33d3173bb679e46d9f338c26f262e",
     "vit_l": "3adcc4315b642a4d2101128f611684e8734c41232a17c648ed1693702a49a622",
-    "vit_b": "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912"
+    "vit_b": "ec2df62732614e57411cdcf32a23ffdf28910380d03139ee0f4fcbe91eb8c912",
+    # preliminary finetuned models
+    "vit_h_lm": "c30a580e6ccaff2f4f0fbaf9cad10cee615a915cdd8c7bc4cb50ea9bdba3fc09",
+    "vit_b_lm": "f2b8676f92a123f6f8ac998818118bd7269a559381ec60af4ac4be5c86024a1b",
 }
 
 
@@ -105,7 +111,13 @@ def get_sam_model(
     """
     checkpoint = _get_checkpoint(model_type, checkpoint_path)
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    sam = sam_model_registry[model_type](checkpoint=checkpoint)
+
+    # Our custom model types have a suffix "_...". This suffix needs to be stripped
+    # before calling sam_model_registry.
+    model_type_ = model_type[:5]
+    assert model_type_ in ("vit_h", "vit_b", "vit_l")
+
+    sam = sam_model_registry[model_type_](checkpoint=checkpoint)
     sam.to(device=device)
     predictor = SamPredictor(sam)
     if return_sam: