Implement CLI for the annotators

Author: constantinpape

micro_sam/sam_annotator/annotator_2d.py

@@ -120,3 +120,59 @@ def clear_prompts(v):
     # clear the initial points needed for workaround
     clear_prompts(v)
     napari.run()
+
+
+def main():
+    import argparse
+    import warnings
+
+    parser = argparse.ArgumentParser(
+        description="Run interactive segmentation for an image."
+    )
+    parser.add_argument(
+        "-i", "--input", required=True,
+        help="The filepath to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "
+        "or elf.io.open_file (e.g. hdf5, zarr, mrc) For the latter you also need to pass the 'key' parameter."
+    )
+    parser.add_argument(
+        "-k", "--key",
+        help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
+        "for a image series it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
+    )
+    parser.add_argument(
+        "-e", "--embedding_path",
+        help="The filepath for saving/loading the pre-computed image embeddings. "
+        "NOTE: It is recommended to pass this argument and store the embeddings, "
+        "otherwise they will be recomputed every time (which can take a long time)."
+    )
+    parser.add_argument(
+        "-s", "--segmentation",
+        help="Optional filepath to a precomputed segmentation. If passed this will be used to initialize the "
+        "'committed_objects' layer. This can be useful if you want to correct an existing segmentation or if you "
+        "have saved intermediate results from the annotator and want to continue with your annotations. "
+        "Supports the same file formats as 'input'."
+    )
+    parser.add_argument(
+        "-sk", "--segmentation_key",
+        help="The key for opening the segmentation data. Same rules as for 'key' apply."
+    )
+    parser.add_argument(
+        "--show_embeddings", action="store_true",
+        help="Visualize the embeddings computed by SegmentAnything. This can be helpful for debugging."
+    )
+
+    args = parser.parse_args()
+    raw = util.load_image_data(args.input, ndim=2, key=args.key)
+
+    if args.segmentation is None:
+        segmentation = None
+    else:
+        segmentation = util.load_image_data(args.segmentation, args.segmentation_key)
+
+    if args.embedding_path is None:
+        warnings.warn("You have not passed an embedding_path. Restarting the annotator may take a long time.")
+
+    annotator_2d(
+        raw, embedding_path=args.embedding_path,
+        show_embeddings=args.show_embeddings, segmentation_result=segmentation
+    )

micro_sam/sam_annotator/annotator_3d.py

@@ -249,3 +249,59 @@ def clear_prompts(v):
     # clear the initial points needed for workaround
     clear_prompts(v)
     napari.run()
+
+
+def main():
+    import argparse
+    import warnings
+
+    parser = argparse.ArgumentParser(
+        description="Run interactive segmentation for an image volume."
+    )
+    parser.add_argument(
+        "-i", "--input", required=True,
+        help="The filepath to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "
+        "or elf.io.open_file (e.g. hdf5, zarr, mrc) For the latter you also need to pass the 'key' parameter."
+    )
+    parser.add_argument(
+        "-k", "--key",
+        help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
+        "for a image series it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
+    )
+    parser.add_argument(
+        "-e", "--embedding_path",
+        help="The filepath for saving/loading the pre-computed image embeddings. "
+        "NOTE: It is recommended to pass this argument and store the embeddings, "
+        "otherwise they will be recomputed every time (which can take a long time)."
+    )
+    parser.add_argument(
+        "-s", "--segmentation",
+        help="Optional filepath to a precomputed segmentation. If passed this will be used to initialize the "
+        "'committed_objects' layer. This can be useful if you want to correct an existing segmentation or if you "
+        "have saved intermediate results from the annotator and want to continue with your annotations. "
+        "Supports the same file formats as 'input'."
+    )
+    parser.add_argument(
+        "-sk", "--segmentation_key",
+        help="The key for opening the segmentation data. Same rules as for 'key' apply."
+    )
+    parser.add_argument(
+        "--show_embeddings", action="store_true",
+        help="Visualize the embeddings computed by SegmentAnything. This can be helpful for debugging."
+    )
+
+    args = parser.parse_args()
+    raw = util.load_image_data(args.input, ndim=3, key=args.key)
+
+    if args.segmentation is None:
+        segmentation = None
+    else:
+        segmentation = util.load_image_data(args.segmentation, args.segmentation_key)
+
+    if args.embedding_path is None:
+        warnings.warn("You have not passed an embedding_path. Restarting the annotator may take a long time.")
+
+    annotator_3d(
+        raw, embedding_path=args.embedding_path,
+        show_embeddings=args.show_embeddings, segmentation_result=segmentation
+    )

micro_sam/sam_annotator/annotator_tracking.py

@@ -199,3 +199,60 @@ def clear_prompts(v):
     # clear the initial points needed for workaround
     clear_prompts(v)
     napari.run()
+
+
+def main():
+    import argparse
+    import warnings
+
+    parser = argparse.ArgumentParser(
+        description="Run interactive segmentation for an image volume."
+    )
+    parser.add_argument(
+        "-i", "--input", required=True,
+        help="The filepath to the image data. Supports all data types that can be read by imageio (e.g. tif, png, ...) "
+        "or elf.io.open_file (e.g. hdf5, zarr, mrc) For the latter you also need to pass the 'key' parameter."
+    )
+    parser.add_argument(
+        "-k", "--key",
+        help="The key for opening data with elf.io.open_file. This is the internal path for a hdf5 or zarr container, "
+        "for a image series it is a wild-card, e.g. '*.png' and for mrc it is 'data'."
+    )
+    parser.add_argument(
+        "-e", "--embedding_path",
+        help="The filepath for saving/loading the pre-computed image embeddings. "
+        "NOTE: It is recommended to pass this argument and store the embeddings, "
+        "otherwise they will be recomputed every time (which can take a long time)."
+    )
+    # Not implemented for the tracking annotator yet.
+    # And we should change the name for it.
+    # parser.add_argument(
+    #     "-s", "--segmentation",
+    #     help="Optional filepath to a precomputed segmentation. If passed this will be used to initialize the "
+    #     "'committed_objects' layer. This can be useful if you want to correct an existing segmentation or if you "
+    #     "have saved intermediate results from the annotator and want to continue with your annotations. "
+    #     "Supports the same file formats as 'input'."
+    # )
+    # parser.add_argument(
+    #     "-sk", "--segmentation_key",
+    #     help="The key for opening the segmentation data. Same rules as for 'key' apply."
+    # )
+    parser.add_argument(
+        "--show_embeddings", action="store_true",
+        help="Visualize the embeddings computed by SegmentAnything. This can be helpful for debugging."
+    )
+
+    args = parser.parse_args()
+    raw = util.load_image_data(args.input, ndim=3, key=args.key)
+
+    # if args.segmentation is None:
+    #     segmentation = None
+    # else:
+    #     segmentation = util.load_image_data(args.segmentation, args.segmentation_key)
+
+    if args.embedding_path is None:
+        warnings.warn("You have not passed an embedding_path. Restarting the annotator may take a long time.")
+
+    annotator_tracking(
+        raw, embedding_path=args.embedding_path, show_embeddings=args.show_embeddings
+    )

micro_sam/util.py

@@ -4,13 +4,19 @@
 import numpy as np
 import requests
 import torch
+import vigra
 import zarr
 
-import vigra
+from elf.io import open_file
 from skimage.measure import regionprops
 
 from segment_anything import sam_model_registry, SamPredictor
 
+try:
+    import imageio.v2 as imageio
+except ImportError:
+    import imageio
+
 try:
     from napari.utils import progress as tqdm
 except ImportError:
@@ -269,10 +275,20 @@ def get_cell_center_coordinates(gt, mode="p"):
     return center_coordinates, bbox_coordinates
 
 
+def load_image_data(path, ndim, key=None, lazy_loading=False):
+    if key is None:
+        image_data = imageio.imread(path) if ndim == 2 else imageio.volread(path)
+    else:
+        with open_file(path, mode="r") as f:
+            image_data = f[key]
+            if not lazy_loading:
+                image_data = image_data[:]
+    return image_data
+
+
 # TODO enable passing options for get_sam
 def main():
     import argparse
-    from elf.io import open_file
 
     parser = argparse.ArgumentParser()
     parser.add_argument("-i", "--input_path", required=True)
@@ -281,7 +297,7 @@ def main():
     args = parser.parse_args()
 
     predictor = get_sam_model()
-    with open_file(args.input_path) as f:
+    with open_file(args.input_path, mode="r") as f:
         data = f[args.key]
         precompute_image_embeddings(predictor, data, save_path=args.output_path)
 

setup.py

@@ -8,13 +8,14 @@
     version="0.0.1",
     description="SegmentAnything For Microscopy",
     author="Anwai Archit, Constantin Pape",
-    author_email="[email protected], [email protected]",
     url="https://user.informatik.uni-goettingen.de/~pape41/",
     packages=["micro_sam"],
-    license="MIT",  # TODO add the license and check that it's fine with SegmentAnything
-    # TODO add entry points for the napari annotator scripts
+    license="MIT",
     entry_points={
         "console_scripts": [
+            "micro_sam.annotator_2d = micro_sam.sam_annotator.annotator_2d:main",
+            "micro_sam.annotator_3d = micro_sam.sam_annotator.annotator_3d:main",
+            "micro_sam.annotator_tracking = micro_sam.sam_annotator.annotator_tracking:main",
             "micro_sam.precompute_embeddings = micro_sam.util:main",
         ]
     }