Merge pull request #4 from computational-cell-analytics/tracking

Further implement tracking functionality

Author: constantinpape

README.md

@@ -56,8 +56,8 @@ pip install -e .
 ## Usage
 
 After the installation the three applications for interactive annotations can be started from the command line or within a python script:
-- **2d annotation**: via the command `micro_sam.annotator_2d` or with the function `micro_sam.sam_annotator.annotator_2d` from python. Run `micro_sam.annotator_2d -h` or check out [examples/sam_annotator_2d](https://github.com/computational-cell-analytics/micro-sam/blob/master/examples/sam_annotator_2d.py) for details. 
-- **3d annotation**: via the command `micro_sam.annotator_3d` or with the function `micro_sam.sam_annotator.annotator_3d` from python. Run `micro_sam.annotator_3d -h` or check out [examples/sam_annotator_3d](https://github.com/computational-cell-analytics/micro-sam/blob/master/examples/sam_annotator_3d.py) for details. 
+- **2d segmentation**: via the command `micro_sam.annotator_2d` or with the function `micro_sam.sam_annotator.annotator_2d` from python. Run `micro_sam.annotator_2d -h` or check out [examples/sam_annotator_2d](https://github.com/computational-cell-analytics/micro-sam/blob/master/examples/sam_annotator_2d.py) for details. 
+- **3d segmentation**: via the command `micro_sam.annotator_3d` or with the function `micro_sam.sam_annotator.annotator_3d` from python. Run `micro_sam.annotator_3d -h` or check out [examples/sam_annotator_3d](https://github.com/computational-cell-analytics/micro-sam/blob/master/examples/sam_annotator_3d.py) for details. 
 - **tracking**: via the command `micro_sam.annotator_tracking` or with the function `micro_sam.sam_annotator.annotator_tracking` from python. Run `micro_sam.annotator_tracking -h` or check out [examples/sam_annotator_tracking](https://github.com/computational-cell-analytics/micro-sam/blob/master/examples/sam_annotator_tracking.py) for details. 
 
 TODO
@@ -94,4 +94,4 @@ micro_sam <- library with utility functionality for using SAM for microscopy dat
 
 If you are using this repository in your research please cite
 - [SegmentAnything](https://arxiv.org/abs/2304.02643)
-- and our repository on [zenodo](TODO) (we are working on a full publication)
+- and our repository on [zenodo](TODO) (we are working on a publication)

development/.gitignore

@@ -0,0 +1,2 @@
+embeddings/
+*.npy

examples/instance_segmentation.py development/instance_segmentation.pyexamples/instance_segmentation.py renamed to development/instance_segmentation.py

@@ -7,12 +7,12 @@
 
 
 def mito_segmentation():
-    input_path = "./data/Lucchi++/Test_In"
+    input_path = "../examples/data/Lucchi++/Test_In"
     with open_file(input_path) as f:
         raw = f["*.png"][-1, :768, :768]
 
     predictor = util.get_sam_model()
-    image_embeddings = util.precompute_image_embeddings(predictor, raw, "./embeddings/embeddings-mito2d.zarr")
+    image_embeddings = util.precompute_image_embeddings(predictor, raw, "../examples/embeddings/embeddings-mito2d.zarr")
     embedding_pca = compute_pca(image_embeddings["features"])
 
     seg, initial_seg = segment_from_embeddings(predictor, image_embeddings=image_embeddings, return_initial_seg=True)
@@ -26,14 +26,16 @@ def mito_segmentation():
 
 
 def cell_segmentation():
-    path = "./DIC-C2DH-HeLa/train/01"
+    path = "../examples/data/DIC-C2DH-HeLa/train/01"
     with open_file(path, mode="r") as f:
         timeseries = f["*.tif"][:50]
 
     frame = 11
 
     predictor = util.get_sam_model()
-    image_embeddings = util.precompute_image_embeddings(predictor, timeseries, "./embeddings/embeddings-ctc.zarr")
+    image_embeddings = util.precompute_image_embeddings(
+        predictor, timeseries, "../examples/embeddings/embeddings-ctc.zarr"
+    )
     embedding_pca = compute_pca(image_embeddings["features"][frame])
 
     seg, initial_seg = segment_from_embeddings(

development/tracking.py

@@ -0,0 +1,46 @@
+from glob import glob
+
+import numpy as np
+from elf.io import open_file
+from micro_sam.sam_annotator import annotator_tracking
+
+
+def debug_tracking(timeseries, embedding_path):
+    import micro_sam.util as util
+    from micro_sam.sam_annotator.annotator_tracking import _track_from_prompts
+
+    predictor = util.get_sam_model()
+    image_embeddings = util.precompute_image_embeddings(predictor, timeseries, embedding_path)
+
+    # seg = np.zeros(timeseries.shape, dtype="uint32")
+    seg = np.load("./seg.npy")
+    assert seg.shape == timeseries.shape
+    slices = np.array([0])
+    stop_upper = False
+
+    _track_from_prompts(seg, predictor, slices, image_embeddings, stop_upper, threshold=0.5, projection="bounding_box")
+
+
+def load_data():
+    pattern = "/home/pape/Work/data/incu_cyte/carmello/videos/MiaPaCa_flat_B3-3_registered/image-*"
+    paths = glob(pattern)
+    paths.sort()
+
+    timeseries = []
+    for p in paths[:45]:
+        with open_file(p, mode="r") as f:
+            timeseries.append(f["phase-contrast"][:])
+    timeseries = np.stack(timeseries)
+    return timeseries
+
+
+def main():
+    timeseries = load_data()
+    embedding_path = "./embeddings/embeddings-tracking.zarr"
+
+    # _check_tracking(timeseries, embedding_path)
+    annotator_tracking(timeseries, embedding_path=embedding_path)
+
+
+if __name__ == "__main__":
+    main()

examples/image_series_annotator_app.py

@@ -11,7 +11,7 @@
 
 from magicgui import magicgui
 from micro_sam.segment_from_prompts import segment_from_points
-from micro_sam.sam_annotator.util import create_prompt_menu, prompt_layer_to_points
+from micro_sam.sam_annotator.util import create_prompt_menu, prompt_layer_to_points, toggle_label
 from napari import Viewer
 
 
@@ -64,15 +64,8 @@ def image_series_annotator(image_paths, embedding_save_path, output_folder):
 
     # toggle the points between positive / negative
     @v.bind_key("t")
-    def toggle_label(event=None):
-        # get the currently selected label
-        current_properties = prompts.current_properties
-        current_label = current_properties["label"][0]
-        new_label = "negative" if current_label == "positive" else "positive"
-        current_properties["label"] = np.array([new_label])
-        prompts.current_properties = current_properties
-        prompts.refresh()
-        prompts.refresh_colors()
+    def _toggle_label(event=None):
+        toggle_label(prompts)
 
     # bind the segmentation to a key 's'
     @v.bind_key("s")

examples/sam_annotator_2d.py

@@ -13,10 +13,11 @@ def livecell_annotator():
 
 def main():
     # 2d annotator for livecell data
-    # livecell_annotator()
+    livecell_annotator()
 
+    # TODO
     # 2d annotator for cell tracking challenge hela data
-    hela_2d_annotator()
+    # hela_2d_annotator()
 
 
 if __name__ == "__main__":

examples/sam_annotator_tracking.py

@@ -1,38 +1,19 @@
-from glob import glob
-
-import numpy as np
 from elf.io import open_file
 from micro_sam.sam_annotator import annotator_tracking
 
 
-def track_incucyte_data():
-    pattern = "/home/pape/Work/data/incu_cyte/carmello/videos/MiaPaCa_flat_B3-3_registered/image-*"
-    paths = glob(pattern)
-    paths.sort()
-
-    timeseries = []
-    for p in paths[:45]:
-        with open_file(p, mode="r") as f:
-            timeseries.append(f["phase-contrast"][:])
-    timeseries = np.stack(timeseries)
-
-    annotator_tracking(timeseries, embedding_path="./embeddings/embeddings-tracking.zarr", show_embeddings=False)
-
-
-# TODO describe how to get the data from CTC
+# This runs the interactive tracking annotator for data from the cell tracking challenge:
+# It uses the training data for the HeLA dataset. You can download the data via
+# http://data.celltrackingchallenge.net/training-datasets/DIC-C2DH-HeLa.zip
 def track_ctc_data():
     path = "./data/DIC-C2DH-HeLa/train/01"
     with open_file(path, mode="r") as f:
         timeseries = f["*.tif"][:50]
-
     annotator_tracking(timeseries, embedding_path="./embeddings/embeddings-ctc.zarr")
 
 
 def main():
-    # private data used for initial tests
-    # track_incucyte_data()
-
-    # data from the cell tracking challenges
+    # run interactive tracking for data from the cell tracking challenge
     track_ctc_data()
 
 

micro_sam/sam_annotator/annotator_2d.py

@@ -8,9 +8,9 @@
 from ..visualization import project_embeddings_for_visualization
 from ..segment_instances import segment_from_embeddings
 from ..segment_from_prompts import segment_from_points
-from .util import commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points
-
-COLOR_CYCLE = ["#00FF00", "#FF0000"]
+from .util import (
+    commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points, toggle_label, LABEL_COLOR_CYCLE
+)
 
 
 @magicgui(call_button="Segment Object [S]")
@@ -35,7 +35,7 @@ def annotator_2d(raw, embedding_path=None, show_embeddings=False, segmentation_r
     global PREDICTOR, IMAGE_EMBEDDINGS
 
     PREDICTOR = util.get_sam_model()
-    IMAGE_EMBEDDINGS = util.precompute_image_embeddings(PREDICTOR, raw, save_path=embedding_path)
+    IMAGE_EMBEDDINGS = util.precompute_image_embeddings(PREDICTOR, raw, save_path=embedding_path, ndim=2)
     util.set_precomputed(PREDICTOR, IMAGE_EMBEDDINGS)
 
     #
@@ -45,16 +45,23 @@ def annotator_2d(raw, embedding_path=None, show_embeddings=False, segmentation_r
     v = Viewer()
 
     v.add_image(raw)
-    v.add_labels(data=np.zeros(raw.shape, dtype="uint32"), name="auto_segmentation")
+    if raw.ndim == 2:
+        shape = raw.shape
+    elif raw.ndim == 3 and raw.shape[-1] == 3:
+        shape = raw.shape[:2]
+    else:
+        raise ValueError(f"Invalid input image of shape {raw.shape}. Expect either 2D grayscale or 3D RGB image.")
+
+    v.add_labels(data=np.zeros(shape, dtype="uint32"), name="auto_segmentation")
     if segmentation_result is None:
-        v.add_labels(data=np.zeros(raw.shape, dtype="uint32"), name="committed_objects")
+        v.add_labels(data=np.zeros(shape, dtype="uint32"), name="committed_objects")
     else:
         v.add_labels(segmentation_result, name="committed_objects")
-    v.add_labels(data=np.zeros(raw.shape, dtype="uint32"), name="current_object")
+    v.add_labels(data=np.zeros(shape, dtype="uint32"), name="current_object")
 
     # show the PCA of the image embeddings
     if show_embeddings:
-        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS["features"], raw.shape)
+        embedding_vis, scale = project_embeddings_for_visualization(IMAGE_EMBEDDINGS["features"], shape)
         v.add_image(embedding_vis, name="embeddings", scale=scale)
 
     labels = ["positive", "negative"]
@@ -63,7 +70,7 @@ def annotator_2d(raw, embedding_path=None, show_embeddings=False, segmentation_r
         name="prompts",
         properties={"label": labels},
         edge_color="label",
-        edge_color_cycle=COLOR_CYCLE,
+        edge_color_cycle=LABEL_COLOR_CYCLE,
         symbol="o",
         face_color="transparent",
         edge_width=0.5,
@@ -98,15 +105,8 @@ def _commit(v):
         commit_segmentation_widget(v)
 
     @v.bind_key("t")
-    def toggle_label(event=None):
-        # get the currently selected label
-        current_properties = prompts.current_properties
-        current_label = current_properties["label"][0]
-        new_label = "negative" if current_label == "positive" else "positive"
-        current_properties["label"] = np.array([new_label])
-        prompts.current_properties = current_properties
-        prompts.refresh()
-        prompts.refresh_colors()
+    def _toggle_label(event=None):
+        toggle_label(prompts)
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):

micro_sam/sam_annotator/annotator_3d.py

@@ -8,9 +8,11 @@
 from .. import util
 from ..segment_from_prompts import segment_from_mask, segment_from_points
 from ..visualization import project_embeddings_for_visualization
-from .util import commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points, segment_slices_with_prompts
-
-COLOR_CYCLE = ["#00FF00", "#FF0000"]
+from .util import (
+    commit_segmentation_widget, create_prompt_menu,
+    prompt_layer_to_points, segment_slices_with_prompts,
+    toggle_label, LABEL_COLOR_CYCLE
+)
 
 
 #
@@ -74,9 +76,16 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
         for z_start, z_stop in zip(segmented_slices[:-1], segmented_slices[1:]):
             slice_diff = z_stop - z_start
             z_mid = int((z_start + z_stop) // 2)
+
             if slice_diff == 1:  # the slices are adjacent -> we don't need to do anything
                 pass
 
+            elif z_start == z0 and stop_lower:  # the lower slice is stop: we just segment from upper
+                segment_range(z_stop, z_start, -1, np.less_equal, verbose=verbose)
+
+            elif z_stop == z1 and stop_upper:  # the upper slice is stop: we just segment from lower
+                segment_range(z_start, z_stop, 1, np.greater_equal, verbose=verbose)
+
             elif slice_diff == 2:  # there is only one slice in between -> use combined mask
                 z = z_start + 1
                 seg_prompt = np.logical_or(seg[z_start] == 1, seg[z_stop] == 1)
@@ -187,7 +196,7 @@ def annotator_3d(raw, embedding_path=None, show_embeddings=False, segmentation_r
         name="prompts",
         properties={"label": labels},
         edge_color="label",
-        edge_color_cycle=COLOR_CYCLE,
+        edge_color_cycle=LABEL_COLOR_CYCLE,
         symbol="o",
         face_color="transparent",
         edge_width=0.5,
@@ -227,15 +236,8 @@ def _commit(v):
         commit_segmentation_widget(v)
 
     @v.bind_key("t")
-    def toggle_label(event=None):
-        # get the currently selected label
-        current_properties = prompts.current_properties
-        current_label = current_properties["label"][0]
-        new_label = "negative" if current_label == "positive" else "positive"
-        current_properties["label"] = np.array([new_label])
-        prompts.current_properties = current_properties
-        prompts.refresh()
-        prompts.refresh_colors()
+    def _toggle_label(event=None):
+        toggle_label(prompts)
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):