Merge pull request #97 from computational-cell-analytics/clean-gui

constantinpape · web-flow · commit 449e025c1ac8 · 2023-07-26T15:14:48.000+02:00
Add gui element to clear annotations
diff --git a/micro_sam/sam_annotator/annotator_2d.py b/micro_sam/sam_annotator/annotator_2d.py
@@ -11,27 +11,23 @@
 from .. import util
 from .. import instance_segmentation
 from ..visualization import project_embeddings_for_visualization
+from . import util as vutil
 from .gui_utils import show_wrong_file_warning
-from .util import (
-    clear_all_prompts, commit_segmentation_widget, create_prompt_menu,
-    prompt_layer_to_boxes, prompt_layer_to_points, prompt_segmentation, toggle_label, LABEL_COLOR_CYCLE,
-    _initialize_parser,
-)
 
 
 @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"])
+    boxes = vutil.prompt_layer_to_boxes(v.layers["box_prompts"])
+    points, labels = vutil.prompt_layer_to_points(v.layers["prompts"])
 
     shape = v.layers["current_object"].data.shape
     if IMAGE_EMBEDDINGS["original_size"] is None:  # tiled prediction
-        seg = prompt_segmentation(
+        seg = vutil.prompt_segmentation(
             PREDICTOR, points, labels, boxes, shape, image_embeddings=IMAGE_EMBEDDINGS, multiple_box_prompts=True
         )
     else:  # normal prediction and we have set the precomputed embeddings already
-        seg = prompt_segmentation(PREDICTOR, points, labels, boxes, shape, multiple_box_prompts=True)
+        seg = vutil.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:
@@ -141,7 +137,7 @@ def _initialize_viewer(raw, segmentation_result, tile_shape, show_embeddings):
         name="prompts",
         properties={"label": labels},
         edge_color="label",
-        edge_color_cycle=LABEL_COLOR_CYCLE,
+        edge_color_cycle=vutil.LABEL_COLOR_CYCLE,
         symbol="o",
         face_color="transparent",
         edge_width=0.5,
@@ -158,14 +154,15 @@ def _initialize_viewer(raw, segmentation_result, tile_shape, show_embeddings):
     # add the widgets
     #
 
-    prompt_widget = create_prompt_menu(prompts, labels)
+    prompt_widget = vutil.create_prompt_menu(prompts, labels)
     v.window.add_dock_widget(prompt_widget)
 
     # (optional) auto-segmentation functionality
     v.window.add_dock_widget(autosegment_widget)
 
     v.window.add_dock_widget(segment_wigdet)
-    v.window.add_dock_widget(commit_segmentation_widget)
+    v.window.add_dock_widget(vutil.commit_segmentation_widget)
+    v.window.add_dock_widget(vutil.clear_widget)
 
     #
     # key bindings
@@ -177,15 +174,15 @@ def _segmet(v):
 
     @v.bind_key("c")
     def _commit(v):
-        commit_segmentation_widget(v)
+        vutil.commit_segmentation_widget(v)
 
     @v.bind_key("t")
     def _toggle_label(event=None):
-        toggle_label(prompts)
+        vutil.toggle_label(prompts)
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        clear_all_prompts(v)
+        vutil.clear_annotations(v)
 
     return v
 
@@ -246,7 +243,7 @@ def annotator_2d(
     #
     # start the viewer
     #
-    clear_all_prompts(v)
+    vutil.clear_annotations(v, clear_segmentations=False)
 
     if return_viewer:
         return v
@@ -255,7 +252,7 @@ def annotator_2d(
 
 
 def main():
-    parser = _initialize_parser(description="Run interactive segmentation for an image.")
+    parser = vutil._initialize_parser(description="Run interactive segmentation for an image.")
     args = parser.parse_args()
     raw = util.load_image_data(args.input, key=args.key)
 
diff --git a/micro_sam/sam_annotator/annotator_3d.py b/micro_sam/sam_annotator/annotator_3d.py
@@ -11,13 +11,8 @@
 from .. import util
 from ..prompt_based_segmentation import segment_from_mask
 from ..visualization import project_embeddings_for_visualization
+from . import util as vutil
 from .gui_utils import show_wrong_file_warning
-from .util import (
-    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,
-    _initialize_parser
-)
 
 
 #
@@ -133,16 +128,16 @@ def segment_slice_wigdet(v: Viewer):
     position = v.cursor.position
     z = int(position[0])
 
-    point_prompts = prompt_layer_to_points(v.layers["prompts"], z)
+    point_prompts = vutil.prompt_layer_to_points(v.layers["prompts"], z)
     # this is a stop prompt, we do nothing
     if not point_prompts:
         return
 
-    boxes = prompt_layer_to_boxes(v.layers["box_prompts"], z)
+    boxes = vutil.prompt_layer_to_boxes(v.layers["box_prompts"], z)
     points, labels = point_prompts
 
     shape = v.layers["current_object"].data.shape[1:]
-    seg = prompt_segmentation(
+    seg = vutil.prompt_segmentation(
         PREDICTOR, points, labels, boxes, shape, multiple_box_prompts=False,
         image_embeddings=IMAGE_EMBEDDINGS, i=z
     )
@@ -175,7 +170,7 @@ def segment_volume_widget(
 
     with progress(total=shape[0]) as progress_bar:
 
-        seg, slices, stop_lower, stop_upper = segment_slices_with_prompts(
+        seg, slices, stop_lower, stop_upper = vutil.segment_slices_with_prompts(
             PREDICTOR, v.layers["prompts"], v.layers["box_prompts"], IMAGE_EMBEDDINGS, shape, progress_bar=progress_bar,
         )
 
@@ -235,7 +230,7 @@ def annotator_3d(
         name="prompts",
         properties={"label": labels},
         edge_color="label",
-        edge_color_cycle=LABEL_COLOR_CYCLE,
+        edge_color_cycle=vutil.LABEL_COLOR_CYCLE,
         symbol="o",
         face_color="transparent",
         edge_width=0.5,
@@ -254,13 +249,14 @@ def annotator_3d(
 
     # TODO add (optional) auto-segmentation functionality
 
-    prompt_widget = create_prompt_menu(prompts, labels)
+    prompt_widget = vutil.create_prompt_menu(prompts, labels)
     v.window.add_dock_widget(prompt_widget)
 
     v.window.add_dock_widget(segment_slice_wigdet)
 
     v.window.add_dock_widget(segment_volume_widget)
-    v.window.add_dock_widget(commit_segmentation_widget)
+    v.window.add_dock_widget(vutil.commit_segmentation_widget)
+    v.window.add_dock_widget(vutil.clear_widget)
 
     #
     # key bindings
@@ -276,30 +272,30 @@ def _seg_volume(v):
 
     @v.bind_key("c")
     def _commit(v):
-        commit_segmentation_widget(v)
+        vutil.commit_segmentation_widget(v)
 
     @v.bind_key("t")
     def _toggle_label(event=None):
-        toggle_label(prompts)
+        vutil.toggle_label(prompts)
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        clear_all_prompts(v)
+        vutil.clear_annotations(v)
 
     #
     # start the viewer
     #
 
     # clear the initial points needed for workaround
-    clear_prompts(v)
+    vutil.clear_annotations(v, clear_segmentations=False)
 
     if return_viewer:
         return v
     napari.run()
 
 
 def main():
-    parser = _initialize_parser(description="Run interactive segmentation for an image volume.")
+    parser = vutil._initialize_parser(description="Run interactive segmentation for an image volume.")
     args = parser.parse_args()
     raw = util.load_image_data(args.input, key=args.key)
 
diff --git a/micro_sam/sam_annotator/annotator_tracking.py b/micro_sam/sam_annotator/annotator_tracking.py
@@ -15,15 +15,9 @@
 
 from .. import util
 from ..prompt_based_segmentation import segment_from_mask
-from .gui_utils import show_wrong_file_warning
-from .util import (
-    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,
-    _initialize_parser
-)
 from ..visualization import project_embeddings_for_visualization
+from . import util as vutil
+from .gui_utils import show_wrong_file_warning
 
 # Cyan (track) and Magenta (division)
 STATE_COLOR_CYCLE = ["#00FFFF", "#FF00FF", ]
@@ -108,7 +102,7 @@ def _update_motion_model(seg, t, t0, motion_model):
             seg_t = seg[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)
+            track_state = vutil.prompt_layer_to_state(point_prompts, t)
 
         # otherwise project the mask (under the motion model) and segment the next slice from the mask
         else:
@@ -189,16 +183,16 @@ def segment_frame_wigdet(v: Viewer):
     position = v.cursor.position
     t = int(position[0])
 
-    point_prompts = prompt_layer_to_points(v.layers["prompts"], t, track_id=CURRENT_TRACK_ID)
+    point_prompts = vutil.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)
+    boxes = vutil.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(
+    seg = vutil.prompt_segmentation(
         PREDICTOR, points, labels, boxes, shape, multiple_box_prompts=False,
         image_embeddings=IMAGE_EMBEDDINGS, i=t
     )
@@ -230,7 +224,7 @@ def track_objet_widget(
 
     with progress(total=shape[0]) as progress_bar:
         # step 1: segment all slices with prompts
-        seg, slices, _, stop_upper = segment_slices_with_prompts(
+        seg, slices, _, stop_upper = vutil.segment_slices_with_prompts(
             PREDICTOR, v.layers["prompts"], v.layers["box_prompts"], IMAGE_EMBEDDINGS, shape,
             progress_bar=progress_bar, track_id=CURRENT_TRACK_ID
         )
@@ -349,7 +343,7 @@ def commit_tracking_widget(v: Viewer, layer: str = "current_track"):
     v.layers[layer].data = np.zeros(shape, dtype="uint32")
     v.layers[layer].refresh()
 
-    clear_all_prompts(v)
+    vutil.clear_annotations(v, clear_segmentations=False)
 
 
 def annotator_tracking(
@@ -409,7 +403,7 @@ def annotator_tracking(
             "track_id": ["1", "1"],  # NOTE we use string to avoid pandas warnings...
         },
         edge_color="label",
-        edge_color_cycle=LABEL_COLOR_CYCLE,
+        edge_color_cycle=vutil.LABEL_COLOR_CYCLE,
         symbol="o",
         face_color="state",
         face_color_cycle=STATE_COLOR_CYCLE,
@@ -444,7 +438,7 @@ def annotator_tracking(
 
     # TODO add (optional) auto-segmentation and tracking functionality
 
-    prompt_widget = create_prompt_menu(prompts, labels)
+    prompt_widget = vutil.create_prompt_menu(prompts, labels)
     v.window.add_dock_widget(prompt_widget)
 
     TRACKING_WIDGET = create_tracking_menu(prompts, box_prompts, state_labels, list(LINEAGE.keys()))
@@ -453,6 +447,7 @@ def annotator_tracking(
     v.window.add_dock_widget(segment_frame_wigdet)
     v.window.add_dock_widget(track_objet_widget)
     v.window.add_dock_widget(commit_tracking_widget)
+    v.window.add_dock_widget(vutil.clear_widget)
 
     #
     # key bindings
@@ -468,15 +463,15 @@ def _track_object(v):
 
     @v.bind_key("t")
     def _toggle_label(event=None):
-        toggle_label(prompts)
+        vutil.toggle_label(prompts)
 
     @v.bind_key("c")
     def _commit(v):
         commit_tracking_widget(v)
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        clear_all_prompts(v)
+        vutil.clear_annotations(v)
 
     #
     # start the viewer
@@ -486,15 +481,15 @@ def clear_prompts(v):
     v.dims.current_step = (0,) + tuple(sh // 2 for sh in raw.shape[1:])
 
     # clear the initial points needed for workaround
-    clear_prompts(v)
+    vutil.clear_annotations(v, clear_segmentations=False)
 
     if return_viewer:
         return v
     napari.run()
 
 
 def main():
-    parser = _initialize_parser(
+    parser = vutil._initialize_parser(
         description="Run interactive segmentation for an image volume.",
         with_segmentation_result=False,
     )
diff --git a/micro_sam/sam_annotator/util.py b/micro_sam/sam_annotator/util.py
@@ -12,14 +12,25 @@
 LABEL_COLOR_CYCLE = ["#00FF00", "#FF0000"]
 
 
-# TODO: also clear the current segmentation or tracking
-# TODO: also add menu for it in the annotator
-def clear_all_prompts(v: Viewer) -> None:
+def clear_annotations(v: Viewer, clear_segmentations=True) -> None:
     v.layers["prompts"].data = []
     v.layers["prompts"].refresh()
     if "box_prompts" in v.layers:
         v.layers["box_prompts"].data = []
         v.layers["box_prompts"].refresh()
+    if not clear_segmentations:
+        return
+    if "current_object" in v.layers:
+        v.layers["current_object"].data = np.zeros(v.layers["current_object"].data.shape, dtype="uint32")
+        v.layers["current_object"].refresh()
+    if "current_track" in v.layers:
+        v.layers["current_track"].data = np.zeros(v.layers["current_track"].data.shape, dtype="uint32")
+        v.layers["current_track"].refresh()
+
+
+@magicgui(call_button="Clear Annotations [Shfit + C]")
+def clear_widget(v: Viewer) -> None:
+    clear_annotations(v)
 
 
 @magicgui(call_button="Commit [C]", layer={"choices": ["current_object", "auto_segmentation"]})
@@ -37,7 +48,7 @@ def commit_segmentation_widget(v: Viewer, layer: str = "current_object") -> None
     v.layers[layer].refresh()
 
     if layer == "current_object":
-        clear_all_prompts(v)
+        clear_annotations(v)
 
 
 def create_prompt_menu(points_layer, labels, menu_name="prompt", label_name="label"):
