Add divison label to tracking annotator WIP

Author: constantinpape

micro_sam/sam_annotator/annotator_2d.py

@@ -7,9 +7,7 @@
 from .. import util
 from ..visualization import compute_pca
 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, LABEL_COLOR_CYCLE
 
 
 @magicgui(call_button="Segment Object [S]")
@@ -53,7 +51,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,

micro_sam/sam_annotator/annotator_3d.py

@@ -8,9 +8,8 @@
 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, 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, LABEL_COLOR_CYCLE)
 
 
 #
@@ -178,7 +177,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,

micro_sam/sam_annotator/annotator_tracking.py

@@ -12,9 +12,10 @@
 from .. import util
 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, segment_slices_with_prompts
+from .util import create_prompt_menu, prompt_layer_to_points, segment_slices_with_prompts, LABEL_COLOR_CYCLE
 
-COLOR_CYCLE = ["#00FF00", "#FF0000"]
+# Magenta and Cyan
+STATE_COLOR_CYCLE = ["#FF00FF", "#00FFFF"]
 
 
 #
@@ -191,21 +192,23 @@ def annotator_tracking(raw, embedding_path=None, show_embeddings=False):
     #
     # add the widgets
     #
-    # TODO add the division labels
     labels = ["positive", "negative"]
+    state_labels = ["division", "track"]
     prompts = v.add_points(
         data=[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],  # FIXME workaround
         name="prompts",
-        properties={"label": labels},
+        properties={"label": labels, "state": state_labels},
         edge_color="label",
-        edge_color_cycle=COLOR_CYCLE,
+        edge_color_cycle=LABEL_COLOR_CYCLE,
         symbol="o",
-        face_color="transparent",
-        edge_width=0.5,  # FIXME workaround
+        face_color="state",
+        face_color_cycle=STATE_COLOR_CYCLE,
+        edge_width=0.4,
         size=12,
         ndim=3,
     )
     prompts.edge_color_mode = "cycle"
+    prompts.face_color_mode = "cycle"
 
     #
     # add the widgets
@@ -216,6 +219,9 @@ def annotator_tracking(raw, embedding_path=None, show_embeddings=False):
     prompt_widget = create_prompt_menu(prompts, labels)
     v.window.add_dock_widget(prompt_widget)
 
+    state_widget = create_prompt_menu(prompts, state_labels, menu_name="state", label_name="state")
+    v.window.add_dock_widget(state_widget)
+
     v.window.add_dock_widget(segment_frame_wigdet)
     v.window.add_dock_widget(track_objet_widget)
 

micro_sam/sam_annotator/util.py

@@ -6,6 +6,9 @@
 
 from ..segment_from_prompts import segment_from_points
 
+# Green and Red
+LABEL_COLOR_CYCLE = ["#00FF00", "#FF0000"]
+
 
 @magicgui(call_button="Commit [C]")
 def commit_segmentation_widget(v: Viewer):
@@ -23,20 +26,20 @@ def commit_segmentation_widget(v: Viewer):
     v.layers["prompts"].refresh()
 
 
-def create_prompt_menu(points_layer, labels):
-    label_menu = ComboBox(label="prompts", choices=labels)
+def create_prompt_menu(points_layer, labels, menu_name="prompt", label_name="label"):
+    label_menu = ComboBox(label=menu_name, choices=labels)
     label_widget = Container(widgets=[label_menu])
 
     def update_label_menu(event):
-        new_label = str(points_layer.current_properties["label"][0])
+        new_label = str(points_layer.current_properties[label_name][0])
         if new_label != label_menu.value:
             label_menu.value = new_label
 
     points_layer.events.current_properties.connect(update_label_menu)
 
     def label_changed(new_label):
         current_properties = points_layer.current_properties
-        current_properties["label"] = np.array([new_label])
+        current_properties[label_name] = np.array([new_label])
         points_layer.current_properties = current_properties
         points_layer.refresh_colors()