Add keypoint selection lock widget

Author: jeylau

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/keypoints.py

@@ -99,16 +99,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: