Refactor segmentation from prompts in annotators and support box prompts in 3d annotator

Author: constantinpape

micro_sam/sam_annotator/annotator_2d.py

@@ -7,67 +7,24 @@
 from .. import util
 from .. import segment_instances
 from ..visualization import project_embeddings_for_visualization
-from ..segment_from_prompts import segment_from_box, segment_from_box_and_points, 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 point prompts
+    # 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"])
-    assert len(points) == len(labels)
-    have_points = len(points) > 0
-
-    # get the current box prompts
-    box_layer = v.layers["box_prompts"]
-    have_boxes = box_layer.nshapes > 0
-
-    # segment only with points
-    if have_points and not have_boxes:
-        seg = segment_from_points(PREDICTOR, points, labels).squeeze()
-
-    # segment only with boxes
-    elif not have_points and have_boxes:
-        shape = v.layers["current_object"].data.shape
-        seg = np.zeros(shape, dtype="uint32")
-
-        seg_id = 1
-        for prompt_id in range(box_layer.nshapes):
-            shape_type = box_layer.shape_type[prompt_id]
-
-            # for now we only support segmentation from rectangles.
-            # supporting other shapes would be possible by casting the shape to a mask
-            # and then segmenting from mask and bounding box.
-            # but for this we need to fix issue with resizing the mask for non-square shapes.
-            if shape_type != "rectangle":
-                print(f"You have provided a {shape_type} shape.")
-                print("We currently only support rectangle shapes for prompts and this prompt will be skipped.")
-                continue
-
-            box = box_layer.data[prompt_id]
-            prompt_box = np.array([box[:, 0].min(), box[:, 1].min(), box[:, 0].max(), box[:, 1].max()])
-            mask = segment_from_box(PREDICTOR, prompt_box).squeeze()
-            seg[mask] = seg_id
-            seg_id += 1
-
-    # segment with points and box (currently only one box supported)
-    elif have_points and have_boxes:
-        if box_layer.nshapes > 1:
-            print("You have provided point prompts and more than one box prompt.")
-            print("This setting is currently not supported.")
-            print("When providing both points and prompts you can only segment one object at a time.")
-            return
-
-        box = box_layer.data[0]
-        prompt_box = np.array([box[:, 0].min(), box[:, 1].min(), box[:, 0].max(), box[:, 1].max()])
-        seg = segment_from_box_and_points(PREDICTOR, prompt_box, points, labels).squeeze()
-
-    # no prompts were given, skip segmentation
-    else:
-        print("You haven't given any prompts.")
-        print("Please provide point and/or box prompts.")
+
+    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
@@ -138,7 +95,7 @@ def annotator_2d(raw, embedding_path=None, show_embeddings=False, segmentation_r
     )
     prompts.edge_color_mode = "cycle"
 
-    box_prompts = v.add_shapes(
+    v.add_shapes(
         face_color="transparent", edge_color="green", edge_width=4, name="box_prompts"
     )
 
@@ -173,10 +130,7 @@ def _toggle_label(event=None):
 
     @v.bind_key("Shift-C")
     def clear_prompts(v):
-        prompts.data = []
-        prompts.refresh()
-        box_prompts.data = []
-        box_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
 )
 
 
@@ -121,14 +121,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 +159,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 +217,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 +257,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/util.py

@@ -4,12 +4,20 @@
 from magicgui.widgets import ComboBox, Container
 from napari import Viewer
 
-from ..segment_from_prompts import segment_from_points
+from ..segment_from_prompts import segment_from_box, segment_from_box_and_points, segment_from_points
 
 # Green and Red
 LABEL_COLOR_CYCLE = ["#00FF00", "#FF0000"]
 
 
+def clear_all_prompts(v):
+    v.layers["prompts"].data = []
+    v.layers["prompts"].refresh()
+    if "box_prompts" in v.layers:
+        v.layers["box_prompts"].data = []
+        v.layers["box_prompts"].refresh()
+
+
 @magicgui(call_button="Commit [C]", layer={"choices": ["current_object", "auto_segmentation"]})
 def commit_segmentation_widget(v: Viewer, layer: str = "current_object"):
     seg = v.layers[layer].data
@@ -25,11 +33,7 @@ def commit_segmentation_widget(v: Viewer, layer: str = "current_object"):
     v.layers[layer].refresh()
 
     if layer == "current_object":
-        v.layers["prompts"].data = []
-        v.layers["prompts"].refresh()
-        if "box_prompts" in v.layers:
-            v.layers["box_prompts"].data = []
-            v.layers["box_prompts"].refresh()
+        clear_all_prompts(v)
 
 
 def create_prompt_menu(points_layer, labels, menu_name="prompt", label_name="label"):
@@ -59,7 +63,8 @@ def prompt_layer_to_points(prompt_layer, i=None, track_id=None):
 
     Arguments:
         prompt_layer: the point layer
-        i [int] - index for the data (required for 3d data)
+        i [int] - index for the data (required for 3d or timeseries data)
+        track_id [int] - id of the current track (required for tracking data)
     """
 
     points = prompt_layer.data
@@ -92,6 +97,50 @@ def prompt_layer_to_points(prompt_layer, i=None, track_id=None):
     return this_points, this_labels
 
 
+def prompt_layer_to_boxes(prompt_layer, i=None, track_id=None):
+    """Extract box prompts for SAM from shape layer.
+
+    Arguments:
+        prompt_layer: the point layer
+        i [int] - index for the data (required for 3d or timeseries data)
+        track_id [int] - id of the current track (required for tracking data)
+    """
+    shape_data = prompt_layer.data
+    shape_types = prompt_layer.shape_type
+    assert len(shape_data) == len(shape_types)
+
+    if i is None:
+        # select all boxes that are rectangles
+        boxes = [data for data, stype in zip(shape_data, shape_types) if stype == "rectangle"]
+    else:
+        # we are currently only supporting rectangle shapes.
+        # other shapes could be supported by providing them as rough mask
+        # (and also providing the corresponding bounding box)
+        # but for this we need to figure out the mask prompts for non-square shapes
+        non_rectangle = [stype != "rectangle" for stype in shape_types]
+        if any(non_rectangle):
+            print(f"You have provided {sum(non_rectangle)} shapes that are not rectangles.")
+            print("We currently do not support these as prompts and they will be ignored.")
+        boxes = [
+            data[:, 1:] for data, stype in zip(shape_data, shape_types)
+            if (stype == "rectangle" and (data[:, 0] == i).all())
+        ]
+
+    # TODO support for track_id
+    # if track_id is not None:
+    #     assert i is not None
+    #     track_ids = np.array(list(map(int, prompt_layer.properties["track_id"])))[mask]
+    #     track_id_mask = track_ids == track_id
+    #     this_labels, this_points = this_labels[track_id_mask], this_points[track_id_mask]
+    # assert len(this_points) == len(this_labels)
+
+    # map to correct box format
+    boxes = [
+        np.array([box[:, 0].min(), box[:, 1].min(), box[:, 0].max(), box[:, 1].max()]) for box in boxes
+    ]
+    return boxes
+
+
 def prompt_layer_to_state(prompt_layer, i):
     """Get the state of the track from the prompt layer.
     Only relevant for annotator_tracking.
@@ -116,27 +165,36 @@ def prompt_layer_to_state(prompt_layer, i):
         return "track"
 
 
-def segment_slices_with_prompts(predictor, prompt_layer, image_embeddings, shape, progress_bar=None, track_id=None):
+def segment_slices_with_prompts(
+    predictor, point_prompts, box_prompts, image_embeddings, shape, progress_bar=None, track_id=None
+):
+    """
+    """
+    assert len(shape) == 3
+    image_shape = shape[1:]
     seg = np.zeros(shape, dtype="uint32")
 
-    z_values = prompt_layer.data[:, 0]
+    z_values = point_prompts.data[:, 0]
+    z_values_boxes = np.concatenate([box[:, 0] for box in box_prompts.data])
+
+    # TODO add track id properties to boxes as well, filter z_values_boxes accordingly
     if track_id is not None:
-        track_ids = np.array(list(map(int, prompt_layer.properties["track_id"])))
+        track_ids = np.array(list(map(int, point_prompts.properties["track_id"])))
         assert len(track_ids) == len(z_values)
         z_values = z_values[track_ids == track_id]
 
-    slices = np.unique(z_values).astype("int")
+    slices = np.unique(np.concatenate([z_values, z_values_boxes])).astype("int")
     stop_lower, stop_upper = False, False
 
     def _update_progress():
         if progress_bar is not None:
             progress_bar.update(1)
 
     for i in slices:
-        prompts_i = prompt_layer_to_points(prompt_layer, i, track_id)
+        points_i = prompt_layer_to_points(point_prompts, i, track_id)
 
         # do we end the segmentation at the outer slices?
-        if prompts_i is None:
+        if points_i is None:
 
             if i == slices[0]:
                 stop_lower = True
@@ -149,14 +207,71 @@ def _update_progress():
             _update_progress()
             continue
 
-        points, labels = prompts_i
-        seg_i = segment_from_points(predictor, points, labels, image_embeddings=image_embeddings, i=i)
+        boxes = prompt_layer_to_boxes(box_prompts, i, track_id)
+        points, labels = points_i
+
+        seg_i = prompt_segmentation(
+            predictor, points, labels, boxes, image_shape, multiple_box_prompts=False,
+            image_embeddings=image_embeddings, i=i
+        )
+        if seg_i is None:
+            print(f"The prompts at slice or frame {i} are invalid and the segmentation was skipped.")
+            print("This will lead to a wrong segmentation across slices or frames.")
+            print(f"Please correct the prompts in {i} and rerun the segmentation.")
+            continue
+
         seg[i] = seg_i
         _update_progress()
 
     return seg, slices, stop_lower, stop_upper
 
 
+def prompt_segmentation(
+    predictor, points, labels, boxes, shape, multiple_box_prompts, image_embeddings=None, i=None
+):
+    """
+    """
+    assert len(points) == len(labels)
+    have_points = len(points) > 0
+    have_boxes = len(boxes) > 0
+
+    # no prompts were given, return None
+    if not have_points and not have_boxes:
+        return
+
+    # box and ppint prompts were given
+    elif have_points and have_boxes:
+        if len(boxes) > 1:
+            print("You have provided point prompts and more than one box prompt.")
+            print("This setting is currently not supported.")
+            print("When providing both points and prompts you can only segment one object at a time.")
+            return
+        seg = segment_from_box_and_points(
+            predictor, boxes[0], points, labels, image_embeddings=image_embeddings, i=i
+        ).squeeze()
+
+    # only point prompts were given
+    elif have_points and not have_boxes:
+        seg = segment_from_points(predictor, points, labels, image_embeddings=image_embeddings, i=i).squeeze()
+
+    # only box prompts were given
+    elif not have_points and have_boxes:
+        seg = np.zeros(shape, dtype="uint32")
+
+        if len(boxes) > 1 and not multiple_box_prompts:
+            print("You have provided more than one box annotation. This is not yet supported in the 3d annotator.")
+            print("You can only segment one object at a time in 3d.")
+            return
+
+        seg_id = 1
+        for box in boxes:
+            mask = segment_from_box(predictor, box, image_embeddings=image_embeddings, i=i).squeeze()
+            seg[mask] = seg_id
+            seg_id += 1
+
+    return seg
+
+
 def toggle_label(prompts):
     # get the currently selected label
     current_properties = prompts.current_properties