Merge pull request #53 from DeepLabCut/lock_kpt

Allow users to lock keypoint selection during LOOP labeling mode

Author: MMathisLab

MANIFEST.in

@@ -1,5 +1,6 @@
 include LICENSE
 include README.md
+include src/napari_deeplabcut/assets/*.svg
 
 recursive-exclude * __pycache__
 recursive-exclude * *.py[co]

README.md

@@ -45,7 +45,7 @@ The easiest way to get started is to drop a folder (typically a folder from with
 
 - `2` and `3`, to easily switch between labeling and selection mode
 - `4`, to enable pan & zoom (which is achieved using the mouse wheel or finger scrolling on the Trackpad)
-- `M`, to cycle through regular (sequential), quick, and cycle annotation mode (see the description [here](https://github.com/DeepLabCut/DeepLabCut-label/blob/ee71b0e15018228c98db3b88769e8a8f4e2c0454/dlclabel/layers.py#L9-L19))
+- `M`, to cycle through regular (sequential), quick, and cycle annotation mode (see the description [here](https://github.com/DeepLabCut/napari-deeplabcut/blob/5a5709dd38868341568d66eab548ae8abf37cd63/src/napari_deeplabcut/keypoints.py#L25-L34))
 - `E`, to enable edge coloring (by default, if using this in refinement GUI mode, points with a confidence lower than 0.6 are marked
 in red)
 - `F`, to toggle between animal and body part color scheme.

src/napari_deeplabcut/_tests/test_keypoints.py

@@ -26,13 +26,6 @@ def test_store(store, fake_keypoints):
     assert store.current_keypoint == kpt
     store.next_keypoint()
 
-    store.smart_reset(event=None)
-    assert store.current_keypoint == kpt
-    assert store.current_label == "kpt_0"
-    assert store.current_id == "animal_0"
-    store.current_id = "animal_1"
-    assert store.current_id == "animal_1"
-
     store._find_first_unlabeled_frame(event=None)
     assert store.current_step == store.n_steps - 1
     # Remove a frame to test whether it is correctly found

src/napari_deeplabcut/_tests/test_widgets.py

@@ -90,6 +90,18 @@ def test_keypoints_dropdown_menu(store):
     assert label_menu.count() == 0
 
 
+def test_keypoints_dropdown_menu_smart_reset(store):
+    widget = _widgets.KeypointsDropdownMenu(store)
+    label_menu = widget.menus['label']
+    label_menu.update_to("kpt_2")
+    widget._locked = True
+    widget.smart_reset(event=None)
+    assert label_menu.currentText() == "kpt_2"
+    widget._locked = False
+    widget.smart_reset(event=None)
+    assert label_menu.currentText() == "kpt_0"
+
+
 def test_color_pair():
     pair = _widgets.LabelPair(color="pink", name="kpt", parent=None)
     assert pair.part_name == "kpt"

src/napari_deeplabcut/_widgets.py

@@ -15,8 +15,8 @@
 from napari.layers.utils import color_manager
 from napari.utils.events import Event
 from napari.utils.history import get_save_history, update_save_history
-from qtpy.QtCore import Qt, QTimer, Signal
-from qtpy.QtGui import QPainter
+from qtpy.QtCore import Qt, QTimer, Signal, QSize
+from qtpy.QtGui import QPainter, QIcon
 from qtpy.QtWidgets import (
     QButtonGroup,
     QCheckBox,
@@ -324,6 +324,11 @@ def _store_crop_coordinates(self, *args):
 
     def _form_dropdown_menus(self, store):
         menu = KeypointsDropdownMenu(store)
+        self.viewer.dims.events.current_step.connect(
+            menu.smart_reset,
+            position="last",
+        )
+        menu.smart_reset(event=None)
         self._menus.append(menu)
         layout = QVBoxLayout()
         layout.addWidget(menu)
@@ -447,11 +452,7 @@ def on_insert(self, event):
             layer.events.query_next_frame.connect(store._advance_step)
             layer.bind_key("Shift-Right", store._find_first_unlabeled_frame)
             layer.bind_key("Shift-Left", store._find_first_unlabeled_frame)
-            self.viewer.dims.events.current_step.connect(
-                store.smart_reset,
-                position="last",
-            )
-            store.smart_reset(event=None)
+
             layer.bind_key("Down", store.next_keypoint, overwrite=True)
             layer.bind_key("Up", store.prev_keypoint, overwrite=True)
             layer.face_color_mode = "cycle"
@@ -548,6 +549,7 @@ def __init__(
         super().__init__(parent)
         self.store = store
         self.store.layer.events.current_properties.connect(self.update_menus)
+        self._locked = False
 
         # Map individuals to their respective bodyparts
         self.id2label = defaultdict(list)
@@ -571,10 +573,24 @@ def __init__(
         layout2 = QVBoxLayout()
         for menu in self.menus.values():
             layout2.addWidget(menu)
+        self.lock_button = QPushButton("Lock selection")
+        self.lock_button.setIcon(QIcon('src/napari_deeplabcut/assets/unlock.svg'))
+        self.lock_button.setIconSize(QSize(24, 24))
+        self.lock_button.clicked.connect(self._lock_current_keypoint)
+        layout2.addWidget(self.lock_button)
         group_box.setLayout(layout2)
         layout1.addWidget(group_box)
         self.setLayout(layout1)
 
+    def _lock_current_keypoint(self):
+        self._locked = not self._locked
+        if self._locked:
+            self.lock_button.setText("Unlock selection")
+            self.lock_button.setIcon(QIcon('src/napari_deeplabcut/assets/lock.svg'))
+        else:
+            self.lock_button.setText("Lock selection")
+            self.lock_button.setIcon(QIcon('src/napari_deeplabcut/assets/unlock.svg'))
+
     def update_menus(self, event):
         keypoint = self.store.current_keypoint
         for attr, menu in self.menus.items():
@@ -589,6 +605,18 @@ def refresh_label_menu(self, text: str):
         menu.blockSignals(False)
         menu.addItems(self.id2label[text])
 
+    def smart_reset(self, event):
+        """Set current keypoint to the first unlabeled one."""
+        if self._locked:
+            return
+        unannotated = ""
+        already_annotated = self.store.annotated_keypoints
+        for keypoint in self.store._keypoints:
+            if keypoint not in already_annotated:
+                unannotated = keypoint
+                break
+        self.store.current_keypoint = unannotated if unannotated else self.store._keypoints[0]
+
 
 def create_dropdown_menu(store, items, attr):
     menu = DropdownMenu(items)

src/napari_deeplabcut/assets/lock.svg

@@ -31,6 +31,9 @@ class LabelMode(CycleEnum):
         annotated point actually moves it to the cursor location.
     LOOP: the first point is placed frame by frame, then it wraps
         to the next label at the end and restart from frame 1, etc.
+        Unless the keypoint selection is locked, the dropdown menu is
+        automatically set to the first unlabeled keypoint of
+        the current frame.
     """
 
     SEQUENTIAL = auto()
@@ -49,7 +52,10 @@ def default(cls):
     "QUICK": "Similar to SEQUENTIAL, but trying to add an already\n"
     "annotated point actually moves it to the cursor location.",
     "LOOP": "The first point is placed frame by frame, then it wraps\n"
-    "to the next label at the end and restart from frame 1, etc.",
+    "to the next label at the end and restart from frame 1, etc.\n"
+    "Unless the keypoint selection is locked, the dropdown menu is\n"
+    "automatically set to the first unlabeled keypoint of\n"
+    "the current frame.",
 }
 
 
@@ -99,16 +105,6 @@ def current_keypoint(self, keypoint: Keypoint):
             current_properties["id"] = np.asarray([keypoint.id])
             self.layer.current_properties = current_properties
 
-    def smart_reset(self, event):
-        """Set current keypoint to the first unlabeled one."""
-        unannotated = ""
-        already_annotated = self.annotated_keypoints
-        for keypoint in self._keypoints:
-            if keypoint not in already_annotated:
-                unannotated = keypoint
-                break
-        self.current_keypoint = unannotated if unannotated else self._keypoints[0]
-
     def next_keypoint(self, *args):
         ind = self._keypoints.index(self.current_keypoint) + 1
         if ind <= len(self._keypoints) - 1: