Merge pull request #22 from computational-cell-analytics/box-prompts

Implement box prompts for all annotators

Author: constantinpape

micro_sam/sam_annotator/annotator_2d.py

@@ -7,17 +7,27 @@
 from .. import util
 from .. import segment_instances
 from ..visualization import project_embeddings_for_visualization
-from ..segment_from_prompts import segment_from_points
 from .util import (
-    commit_segmentation_widget, create_prompt_menu, prompt_layer_to_points, toggle_label, LABEL_COLOR_CYCLE
+    clear_all_prompts, commit_segmentation_widget, create_prompt_menu,
+    prompt_layer_to_boxes, prompt_layer_to_points, prompt_segmentation, toggle_label, LABEL_COLOR_CYCLE
 )
 
 
 @magicgui(call_button="Segment Object [S]")
 def segment_wigdet(v: Viewer):
+    # get the current box and point prompts
+    boxes = prompt_layer_to_boxes(v.layers["box_prompts"])
     points, labels = prompt_layer_to_points(v.layers["prompts"])
-    seg = segment_from_points(PREDICTOR, points, labels)
-    v.layers["current_object"].data = seg.squeeze()
+
+    shape = v.layers["current_object"].data.shape
+    seg = prompt_segmentation(PREDICTOR, points, labels, boxes, shape, multiple_box_prompts=True)
+
+    # no prompts were given or prompts were invalid, skip segmentation
+    if seg is None:
+        print("You either haven't provided any prompts or invalid prompts. The segmentation will be skipped.")
+        return
+
+    v.layers["current_object"].data = seg
     v.layers["current_object"].refresh()
 
 
@@ -85,6 +95,10 @@ def annotator_2d(raw, embedding_path=None, show_embeddings=False, segmentation_r
     )
     prompts.edge_color_mode = "cycle"
 
+    v.add_shapes(
+        face_color="transparent", edge_color="green", edge_width=4, name="box_prompts"
+    )
+
     #
     # add the widgets
     #
@@ -116,8 +130,7 @@ def _toggle_label(event=None):
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        prompts.data = []
-        prompts.refresh()
+        clear_all_prompts(v)
 
     #
     # start the viewer

micro_sam/sam_annotator/annotator_3d.py

@@ -6,12 +6,12 @@
 from napari.utils import progress
 
 from .. import util
-from ..segment_from_prompts import segment_from_mask, segment_from_points
+from ..segment_from_prompts import segment_from_mask
 from ..visualization import project_embeddings_for_visualization
 from .util import (
-    commit_segmentation_widget, create_prompt_menu,
-    prompt_layer_to_points, segment_slices_with_prompts,
-    toggle_label, LABEL_COLOR_CYCLE
+    clear_all_prompts, commit_segmentation_widget, create_prompt_menu,
+    prompt_layer_to_boxes, prompt_layer_to_points, prompt_segmentation,
+    segment_slices_with_prompts, toggle_label, LABEL_COLOR_CYCLE
 )
 
 
@@ -95,7 +95,9 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
 
             else:  # there is a range of more than 2 slices in between -> segment ranges
                 # segment from bottom
-                segment_range(z_start, z_mid, 1, np.greater_equal, verbose=verbose)
+                segment_range(
+                    z_start, z_mid, 1, np.greater_equal if slice_diff % 2 == 0 else np.greater, verbose=verbose
+                )
                 # segment from top
                 segment_range(z_stop, z_mid, -1, np.less_equal, verbose=verbose)
                 # if the difference between start and stop is even,
@@ -121,14 +123,26 @@ def segment_slice_wigdet(v: Viewer):
     position = v.cursor.position
     z = int(position[0])
 
-    this_prompts = prompt_layer_to_points(v.layers["prompts"], z)
-    if this_prompts is None:
+    point_prompts = prompt_layer_to_points(v.layers["prompts"], z)
+    # this is a stop prompt, we do nothing
+    if not point_prompts:
         return
 
-    points, labels = this_prompts
-    seg = segment_from_points(PREDICTOR, points, labels, image_embeddings=IMAGE_EMBEDDINGS, i=z)
+    boxes = prompt_layer_to_boxes(v.layers["box_prompts"], z)
+    points, labels = point_prompts
 
-    v.layers["current_object"].data[z] = seg.squeeze()
+    shape = v.layers["current_object"].data.shape[1:]
+    seg = prompt_segmentation(
+        PREDICTOR, points, labels, boxes, shape, multiple_box_prompts=False,
+        image_embeddings=IMAGE_EMBEDDINGS, i=z
+    )
+
+    # no prompts were given or prompts were invalid, skip segmentation
+    if seg is None:
+        print("You either haven't provided any prompts or invalid prompts. The segmentation will be skipped.")
+        return
+
+    v.layers["current_object"].data[z] = seg
     v.layers["current_object"].refresh()
 
 
@@ -147,7 +161,7 @@ def segment_volume_widget(v: Viewer, iou_threshold: float = 0.8, projection: str
     with progress(total=shape[0]) as progress_bar:
 
         seg, slices, stop_lower, stop_upper = segment_slices_with_prompts(
-            PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape, progress_bar=progress_bar,
+            PREDICTOR, v.layers["prompts"], v.layers["box_prompts"], IMAGE_EMBEDDINGS, shape, progress_bar=progress_bar,
         )
 
         # step 2: segment the rest of the volume based on smart prompting
@@ -205,6 +219,10 @@ def annotator_3d(raw, embedding_path=None, show_embeddings=False, segmentation_r
     )
     prompts.edge_color_mode = "cycle"
 
+    v.add_shapes(
+        face_color="transparent", edge_color="green", edge_width=4, name="box_prompts", ndim=3
+    )
+
     #
     # add the widgets
     #
@@ -241,8 +259,7 @@ def _toggle_label(event=None):
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        prompts.data = []
-        prompts.refresh()
+        clear_all_prompts(v)
 
     #
     # start the viewer

micro_sam/sam_annotator/annotator_tracking.py

@@ -11,15 +11,17 @@
 # from vigra.filters import eccentricityCenters
 
 from .. import util
-from ..segment_from_prompts import segment_from_mask, segment_from_points
+from ..segment_from_prompts import segment_from_mask
 from .util import (
-    create_prompt_menu, prompt_layer_to_points, prompt_layer_to_state,
+    create_prompt_menu, clear_all_prompts,
+    prompt_layer_to_boxes, prompt_layer_to_points,
+    prompt_layer_to_state, prompt_segmentation,
     segment_slices_with_prompts, toggle_label, LABEL_COLOR_CYCLE
 )
 from ..visualization import project_embeddings_for_visualization
 
-# Magenta and Cyan
-STATE_COLOR_CYCLE = ["#FF00FF", "#00FFFF"]
+# Cyan (track) and Magenta (division)
+STATE_COLOR_CYCLE = ["#00FFFF", "#FF00FF", ]
 
 
 #
@@ -56,7 +58,7 @@ def _shift_object(mask, motion_model):
 
 # TODO division classifier
 def _track_from_prompts(
-    prompt_layer, seg, predictor, slices, image_embeddings,
+    point_prompts, box_prompts, seg, predictor, slices, image_embeddings,
     stop_upper, threshold, projection,
     progress_bar=None, motion_smoothing=0.5,
 ):
@@ -99,7 +101,9 @@ def _update_motion_model(seg, t, t0, motion_model):
         if t in slices:
             seg_prev = None
             seg_t = seg[t]
-            track_state = prompt_layer_to_state(prompt_layer, t)
+            # currently using the box layer doesn't work for keeping track of the track state
+            # track_state = prompt_layers_to_state(point_prompts, box_prompts, t)
+            track_state = prompt_layer_to_state(point_prompts, t)
 
         # otherwise project the mask (under the motion model) and segment the next slice from the mask
         else:
@@ -138,7 +142,7 @@ def _update_motion_model(seg, t, t0, motion_model):
             break
 
         # stop if we are at the last slce
-        if t == seg.shape[0] - 1:
+        if t == seg.shape[0]:
             break
 
         # stop if we have a division
@@ -180,9 +184,24 @@ def segment_frame_wigdet(v: Viewer):
     position = v.cursor.position
     t = int(position[0])
 
-    this_prompts = prompt_layer_to_points(v.layers["prompts"], t, track_id=CURRENT_TRACK_ID)
-    points, labels = this_prompts
-    seg = segment_from_points(PREDICTOR, points, labels, image_embeddings=IMAGE_EMBEDDINGS, i=t)
+    point_prompts = prompt_layer_to_points(v.layers["prompts"], t, track_id=CURRENT_TRACK_ID)
+    # this is a stop prompt, we do nothing
+    if not point_prompts:
+        return
+
+    boxes = prompt_layer_to_boxes(v.layers["box_prompts"], t, track_id=CURRENT_TRACK_ID)
+    points, labels = point_prompts
+
+    shape = v.layers["current_track"].data.shape[1:]
+    seg = prompt_segmentation(
+        PREDICTOR, points, labels, boxes, shape, multiple_box_prompts=False,
+        image_embeddings=IMAGE_EMBEDDINGS, i=t
+    )
+
+    # no prompts were given or prompts were invalid, skip segmentation
+    if seg is None:
+        print("You either haven't provided any prompts or invalid prompts. The segmentation will be skipped.")
+        return
 
     # clear the old segmentation for this track_id
     old_mask = v.layers["current_track"].data[t] == CURRENT_TRACK_ID
@@ -199,24 +218,23 @@ def track_objet_widget(
 ):
     shape = v.layers["raw"].data.shape
 
-    # choose mask projection for square images and bounding box projection otherwise
-    # (because mask projection does not work properly for non-square images yet)
-    if projection == "default":
-        projection_ = "mask" if shape[1] == shape[2] else "bounding_box"
-    else:
-        projection_ = projection
+    # we use the bounding box projection method as default which generally seems to work better for larger changes
+    # between frames (which is pretty tyipical for tracking compared to 3d segmentation)
+    projection_ = "bounding_box" if projection == "default" else projection
 
     with progress(total=shape[0]) as progress_bar:
         # step 1: segment all slices with prompts
         seg, slices, _, stop_upper = segment_slices_with_prompts(
-            PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape,
+            PREDICTOR, v.layers["prompts"], v.layers["box_prompts"], IMAGE_EMBEDDINGS, shape,
             progress_bar=progress_bar, track_id=CURRENT_TRACK_ID
         )
 
         # step 2: track the object starting from the lowest annotated slice
         seg, has_division = _track_from_prompts(
-            v.layers["prompts"], seg, PREDICTOR, slices, IMAGE_EMBEDDINGS, stop_upper, threshold=iou_threshold,
-            projection=projection_, progress_bar=progress_bar, motion_smoothing=motion_smoothing,
+            v.layers["prompts"], v.layers["box_prompts"], seg,
+            PREDICTOR, slices, IMAGE_EMBEDDINGS, stop_upper,
+            threshold=iou_threshold, projection=projection_,
+            progress_bar=progress_bar, motion_smoothing=motion_smoothing,
         )
 
     # if a division has occurred and it's the first time it occurred for this track
@@ -231,7 +249,7 @@ def track_objet_widget(
     v.layers["current_track"].refresh()
 
 
-def create_tracking_menu(points_layer, states, track_ids):
+def create_tracking_menu(points_layer, box_layer, states, track_ids):
     state_menu = ComboBox(label="track_state", choices=states)
     track_id_menu = ComboBox(label="track_id", choices=list(map(str, track_ids)))
     tracking_widget = Container(widgets=[state_menu, track_id_menu])
@@ -245,11 +263,25 @@ def update_track_id(event):
         global CURRENT_TRACK_ID
         new_id = str(points_layer.current_properties["track_id"][0])
         if new_id != track_id_menu.value:
-            state_menu.value = new_id
+            track_id_menu.value = new_id
+            CURRENT_TRACK_ID = int(new_id)
+
+    # def update_state_boxes(event):
+    #     new_state = str(box_layer.current_properties["state"][0])
+    #     if new_state != state_menu.value:
+    #         state_menu.value = new_state
+
+    def update_track_id_boxes(event):
+        global CURRENT_TRACK_ID
+        new_id = str(box_layer.current_properties["track_id"][0])
+        if new_id != track_id_menu.value:
+            track_id_menu.value = new_id
             CURRENT_TRACK_ID = int(new_id)
 
     points_layer.events.current_properties.connect(update_state)
     points_layer.events.current_properties.connect(update_track_id)
+    # box_layer.events.current_properties.connect(update_state_boxes)
+    box_layer.events.current_properties.connect(update_track_id_boxes)
 
     def state_changed(new_state):
         current_properties = points_layer.current_properties
@@ -264,8 +296,23 @@ def track_id_changed(new_track_id):
         points_layer.current_properties = current_properties
         CURRENT_TRACK_ID = int(new_track_id)
 
+    # def state_changed_boxes(new_state):
+    #     current_properties = box_layer.current_properties
+    #     current_properties["state"] = np.array([new_state])
+    #     box_layer.current_properties = current_properties
+    #     box_layer.refresh_colors()
+
+    def track_id_changed_boxes(new_track_id):
+        global CURRENT_TRACK_ID
+        current_properties = box_layer.current_properties
+        current_properties["track_id"] = np.array([new_track_id])
+        box_layer.current_properties = current_properties
+        CURRENT_TRACK_ID = int(new_track_id)
+
     state_menu.changed.connect(state_changed)
     track_id_menu.changed.connect(track_id_changed)
+    # state_menu.changed.connect(state_changed_boxes)
+    track_id_menu.changed.connect(track_id_changed_boxes)
 
     state_menu.set_choice("track")
     return tracking_widget
@@ -295,8 +342,7 @@ def commit_tracking_widget(v: Viewer, layer: str = "current_track"):
     v.layers[layer].data = np.zeros(shape, dtype="uint32")
     v.layers[layer].refresh()
 
-    v.layers["prompts"].data = []
-    v.layers["prompts"].refresh()
+    clear_all_prompts(v)
 
 
 def annotator_tracking(raw, embedding_path=None, show_embeddings=False, tracking_result=None, model_type="vit_h"):
@@ -333,7 +379,7 @@ def annotator_tracking(raw, embedding_path=None, show_embeddings=False, tracking
     # add the widgets
     #
     labels = ["positive", "negative"]
-    state_labels = ["division", "track"]
+    state_labels = ["track", "division"]
     prompts = v.add_points(
         data=[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],  # FIXME workaround
         name="prompts",
@@ -354,6 +400,24 @@ def annotator_tracking(raw, embedding_path=None, show_embeddings=False, tracking
     prompts.edge_color_mode = "cycle"
     prompts.face_color_mode = "cycle"
 
+    # using the box layer to set divisions currently doesn't work
+    # (and setting new track ids also doesn't work, but keeping track of them in the properties is working)
+    box_prompts = v.add_shapes(
+        data=[
+            np.array([[0, 0, 0], [0, 0, 10], [0, 10, 0], [0, 10, 10]]),
+            np.array([[0, 0, 0], [0, 0, 11], [0, 11, 0], [0, 11, 11]]),
+        ],  # FIXME workaround
+        shape_type="rectangle",  # FIXME workaround
+        edge_width=4, ndim=3,
+        face_color="transparent",
+        name="box_prompts",
+        edge_color="green",
+        properties={"track_id": ["1", "1"]},
+        # properties={"track_id": ["1", "1"], "state": state_labels},
+        # edge_color_cycle=STATE_COLOR_CYCLE,
+    )
+    # box_prompts.edge_color_mode = "cycle"
+
     #
     # add the widgets
     #
@@ -363,7 +427,7 @@ def annotator_tracking(raw, embedding_path=None, show_embeddings=False, tracking
     prompt_widget = create_prompt_menu(prompts, labels)
     v.window.add_dock_widget(prompt_widget)
 
-    TRACKING_WIDGET = create_tracking_menu(prompts, state_labels, list(LINEAGE.keys()))
+    TRACKING_WIDGET = create_tracking_menu(prompts, box_prompts, state_labels, list(LINEAGE.keys()))
     v.window.add_dock_widget(TRACKING_WIDGET)
 
     v.window.add_dock_widget(segment_frame_wigdet)
@@ -392,8 +456,7 @@ def _commit(v):
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        prompts.data = []
-        prompts.refresh()
+        clear_all_prompts(v)
 
     #
     # start the viewer