Implement rudimentary functionality for tracking annotator

Author: constantinpape

README.md

@@ -7,6 +7,7 @@ We implement napari applications for:
 - interactive tracking of 2d image data
 
 **Early beta version**
+
 This is an early beta version. Any feedback is welcome, but please be aware that the functionality is evolving fast and not fully tested.
 
 ## Functionality overview

micro_sam/sam_annotator/annotator_3d.py

@@ -7,7 +7,7 @@
 from .. import util
 from ..segment_from_prompts import segment_from_mask, segment_from_points
 from ..visualization import compute_pca
-from .util import commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points
+from .util import commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points, segment_slices_with_prompts
 
 COLOR_CYCLE = ["#00FF00", "#FF0000"]
 
@@ -18,33 +18,6 @@
 #
 
 
-def _segment_slices_with_prompts(predictor, prompt_layer, image_embeddings, shape):
-    seg = np.zeros(shape, dtype="uint32")
-
-    slices = np.unique(prompt_layer.data[:, 0]).astype("int")
-    stop_lower, stop_upper = False, False
-
-    for z in slices:
-        prompts_z = prompt_layer_to_points(prompt_layer, z)
-
-        # do we end the segmentation at the outer slices?
-        if prompts_z is None:
-            if z == slices[0]:
-                stop_lower = True
-            elif z == slices[-1]:
-                stop_upper = True
-            else:
-                raise RuntimeError("Stop slices can only be at the start or end")
-            seg[z] = 0
-            continue
-
-        points, labels = prompts_z
-        seg_z = segment_from_points(predictor, points, labels, image_embeddings=image_embeddings, i=z)
-        seg[z] = seg_z
-
-    return seg, slices, stop_lower, stop_upper
-
-
 # TODO refactor
 def _segment_volume(
     seg, predictor, image_embeddings, segmented_slices,
@@ -74,9 +47,9 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
             seg[z] = seg_z
             z += increment
             if stopping_criterion(z, z_stop):
-                break
                 if verbose:
                     print(f"Segment {z_start} to {z_stop}: stop at slice {z}")
+                break
 
     z0, z1 = int(segmented_slices.min()), int(segmented_slices.max())
 
@@ -145,7 +118,7 @@ def segment_slice_wigdet(v: Viewer):
 def segment_volume_widget(v: Viewer, iou_threshold: float = 0.8, method: str = "mask"):
     # step 1: segment all slices with prompts
     shape = v.layers["raw"].data.shape
-    seg, slices, stop_lower, stop_upper = _segment_slices_with_prompts(
+    seg, slices, stop_lower, stop_upper = segment_slices_with_prompts(
         PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape
     )
 

micro_sam/sam_annotator/annotator_tracking.py

@@ -5,13 +5,60 @@
 from napari import Viewer
 
 from .. import util
-from ..segment_from_prompts import segment_from_points
+from ..segment_from_prompts import segment_from_mask, segment_from_points
 from ..visualization import compute_pca
-from .util import create_prompt_menu, prompt_layer_to_points
+from .util import create_prompt_menu, prompt_layer_to_points, segment_slices_with_prompts
 
 COLOR_CYCLE = ["#00FF00", "#FF0000"]
 
 
+#
+# util functionality
+#
+
+
+# TODO motion model!!!
+# TODO handle divison annotations + division classifier
+def _track_from_prompts(seg, predictor, slices, image_embeddings, stop_upper, threshold, method):
+    assert method in ("mask", "bounding_box")
+    if method == "mask":
+        use_mask, use_box = True, True
+    else:
+        use_mask, use_box = False, True
+
+    t0 = int(slices.min())
+    t = t0 + 1
+    while True:
+        if t in slices:
+            seg_prev = None
+            seg_t = seg[t]
+        else:
+            seg_prev = seg[t - 1]
+            seg_t = segment_from_mask(predictor, seg_prev, image_embeddings=image_embeddings, i=t,
+                                      use_mask=use_mask, use_box=use_box)
+
+        if (threshold is not None) and (seg_prev is not None):
+            iou = util.compute_iou(seg_prev, seg_t)
+            if iou < threshold:
+                msg = f"Segmentation stopped at frame {t} due to IOU {iou} < {threshold}."
+                print(msg)
+                break
+
+        seg[t] = seg_t
+        t += 1
+
+        # TODO here we need to also stop once divisions are implemented
+        # stop tracking if we have stop upper set (i.e. single negative point was set to indicate stop track)
+        if t == slices[-1] and stop_upper:
+            break
+
+        # stop if we are at the last slce
+        if t == seg.shape[0] - 1:
+            break
+
+    return seg
+
+
 #
 # the widgets
 #
@@ -31,8 +78,18 @@ def segment_frame_wigdet(v: Viewer):
 
 
 @magicgui(call_button="Track Object [V]", method={"choices": ["bounding_box", "mask"]})
-def track_objet_widget(v: Viewer, iou_threshold: float = 0.8, method: str = "bounding_box"):
-    pass
+def track_objet_widget(v: Viewer, iou_threshold: float = 0.8, method: str = "mask"):
+    # step 1: segment all slices with prompts
+    shape = v.layers["raw"].data.shape
+    seg, slices, _, stop_upper = segment_slices_with_prompts(
+        PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape
+    )
+
+    # step 2: track the object starting from the lowest annotated slice
+    seg = _track_from_prompts(seg, PREDICTOR, slices, IMAGE_EMBEDDINGS, stop_upper, iou_threshold, method)
+
+    v.layers["current_track"].data = seg
+    v.layers["current_track"].refresh()
 
 
 def annotator_tracking(raw, embedding_path=None, show_embeddings=False):

micro_sam/sam_annotator/util.py

@@ -4,6 +4,8 @@
 from magicgui.widgets import ComboBox, Container
 from napari import Viewer
 
+from ..segment_from_prompts import segment_from_points
+
 
 @magicgui(call_button="Commit [C]")
 def commit_segmentation_widget(v: Viewer):
@@ -72,3 +74,31 @@ def prompt_layer_to_points(prompt_layer, i=None):
         return None
 
     return this_points, this_labels
+
+
+def segment_slices_with_prompts(predictor, prompt_layer, image_embeddings, shape):
+    seg = np.zeros(shape, dtype="uint32")
+
+    slices = np.unique(prompt_layer.data[:, 0]).astype("int")
+    stop_lower, stop_upper = False, False
+
+    for i in slices:
+        prompts_i = prompt_layer_to_points(prompt_layer, i)
+
+        # TODO also take into account division properties once we have this implemented in tracking
+        # do we end the segmentation at the outer slices?
+        if prompts_i is None:
+            if i == slices[0]:
+                stop_lower = True
+            elif i == slices[-1]:
+                stop_upper = True
+            else:
+                raise RuntimeError("Stop slices can only be at the start or end")
+            seg[i] = 0
+            continue
+
+        points, labels = prompts_i
+        seg_i = segment_from_points(predictor, points, labels, image_embeddings=image_embeddings, i=i)
+        seg[i] = seg_i
+
+    return seg, slices, stop_lower, stop_upper