Merge pull request #51 from DeepLabCut/keypoint_diff

Support the addition of new body parts after having started labeling

Author: jeylau

README.md

@@ -101,6 +101,8 @@ Suggested workflows, depending on the image folder contents:
 
     Saving works as described in *1*.
 
+    ***Note that if a new body part has been added to the `config.yaml` file after having started to label, loading the config in the GUI is necessary to update the dropdown menus and other metadata.***
+
 3. **Refining labels** – the image folder contains a `machinelabels-iter<#>.h5` file.
 
     The process is analog to *2*.

src/napari_deeplabcut/_widgets.py

@@ -203,8 +203,8 @@ def __init__(self, napari_viewer):
 
         self._radio_group = self._form_mode_radio_buttons()
 
+        self._display = ColorSchemeDisplay(parent=self)
         self._color_scheme_display = self._form_color_scheme_display(self.viewer)
-
         self._view_scheme_cb.toggled.connect(self._show_color_scheme)
         self._view_scheme_cb.toggle()
 
@@ -354,29 +354,23 @@ def _func():
         return group
 
     def _form_color_scheme_display(self, viewer):
-        display = ColorSchemeDisplay(parent=self)
-        self._update_color_scheme(display)
-
-        self.viewer.layers.events.inserted.connect(
-            partial(self._update_color_scheme, display)
-        )
-
+        self.viewer.layers.events.inserted.connect(self._update_color_scheme)
         return viewer.window.add_dock_widget(
-            display, name="Color scheme reference", area="left"
+            self._display, name="Color scheme reference", area="left"
         )
 
-    def _update_color_scheme(self, display):
+    def _update_color_scheme(self):
+        def to_hex(nparray):
+            a = np.array(nparray * 255, dtype=int)
+            rgb2hex = lambda r, g, b, _: f"#{r:02x}{g:02x}{b:02x}"
+            res = rgb2hex(*a)
+            return res
+
+        self._display.reset()
         for layer in self.viewer.layers:
             if isinstance(layer, Points) and layer.metadata:
-
-                def to_hex(nparray):
-                    a = np.array(nparray * 255, dtype=int)
-                    rgb2hex = lambda r, g, b, _: f"#{r:02x}{g:02x}{b:02x}"
-                    res = rgb2hex(*a)
-                    return res
-
                 [
-                    display.add_entry(name, to_hex(color))
+                    self._display.add_entry(name, to_hex(color))
                     for name, color in layer.metadata["face_color_cycles"][
                         "label"
                     ].items()
@@ -439,6 +433,38 @@ def on_insert(self, event):
                 10, partial(self._move_image_layer_to_bottom, event.index)
             )
         elif isinstance(layer, Points):
+            # If the current Points layer comes from a config file, some have already
+            # been added and the body part names are different from the existing ones,
+            # then we update store's metadata and menus.
+            if layer.metadata.get("project", "") and self._stores:
+                keypoints_menu = self._menus[0].menus["label"]
+                current_keypoint_set = set(
+                    keypoints_menu.itemText(i) for i in range(keypoints_menu.count())
+                )
+                new_keypoint_set = set(layer.metadata["header"].bodyparts)
+                diff = new_keypoint_set.difference(current_keypoint_set)
+                if diff:
+                    answer = QMessageBox.question(self, "", "Do you want to display the new keypoints only?")
+                    if answer  == QMessageBox.Yes:
+                        self.viewer.layers[-2].shown = False
+
+                    self.viewer.status = f"New keypoint{'s' if len(diff) > 1 else ''} {', '.join(diff)} found."
+                    for _layer, store in self._stores.items():
+                        _layer.metadata["header"] = layer.metadata["header"]
+                        _layer.metadata["face_color_cycles"] = layer.metadata["face_color_cycles"]
+                        _layer.face_color_cycle = layer.face_color_cycle
+                        store.layer = _layer
+
+                    for menu in self._menus:
+                        menu._map_individuals_to_bodyparts()
+                        menu._update_items()
+
+                    self._update_color_scheme()
+
+                # Remove the unnecessary layer newly added
+                QTimer.singleShot(10, self.viewer.layers.pop)
+                return
+
             store = keypoints.KeypointStore(self.viewer, layer)
             self._stores[layer] = store
             # TODO Set default dir of the save file dialog
@@ -470,14 +496,15 @@ def on_insert(self, event):
 
     def on_remove(self, event):
         layer = event.value
-        if isinstance(layer, Points):
+        n_points_layer = sum(isinstance(l, Points) for l in self.viewer.layers)
+        if isinstance(layer, Points) and n_points_layer == 0:
             if self._color_scheme_display is not None:
-                self.viewer.window.remove_dock_widget(self._color_scheme_display)
+                self._display.reset()
             self._stores.pop(layer, None)
             while self._menus:
                 menu = self._menus.pop()
                 self._layout.removeWidget(menu)
-                menu.setParent(None)
+                menu.deleteLater()
                 menu.destroy()
             self._trail_cb.setEnabled(False)
             self.last_saved_label.hide()
@@ -529,7 +556,7 @@ def toggle_edge_color(layer):
 class DropdownMenu(QComboBox):
     def __init__(self, labels: Sequence[str], parent: Optional[QWidget] = None):
         super().__init__(parent)
-        self.addItems(labels)
+        self.update_items(labels)
 
     def update_to(self, text: str):
         index = self.findText(text)
@@ -539,6 +566,10 @@ def update_to(self, text: str):
     def reset(self):
         self.setCurrentIndex(0)
 
+    def update_items(self, items):
+        self.clear()
+        self.addItems(items)
+
 
 class KeypointsDropdownMenu(QWidget):
     def __init__(
@@ -551,22 +582,11 @@ def __init__(
         self.store.layer.events.current_properties.connect(self.update_menus)
         self._locked = False
 
-        # Map individuals to their respective bodyparts
         self.id2label = defaultdict(list)
-        for keypoint in store._keypoints:
-            label = keypoint.label
-            id_ = keypoint.id
-            if label not in self.id2label[id_]:
-                self.id2label[id_].append(label)
-
         self.menus = dict()
-        if store.ids[0]:
-            menu = create_dropdown_menu(store, list(self.id2label), "id")
-            menu.currentTextChanged.connect(self.refresh_label_menu)
-            self.menus["id"] = menu
-        self.menus["label"] = create_dropdown_menu(
-            store, self.id2label[store.ids[0]], "label"
-        )
+        self._map_individuals_to_bodyparts()
+        self._populate_menus()
+
         layout1 = QVBoxLayout()
         layout1.addStretch(1)
         group_box = QGroupBox("Keypoint selection")
@@ -582,6 +602,29 @@ def __init__(
         layout1.addWidget(group_box)
         self.setLayout(layout1)
 
+    def _map_individuals_to_bodyparts(self):
+        for keypoint in self.store._keypoints:
+            label = keypoint.label
+            id_ = keypoint.id
+            if label not in self.id2label[id_]:
+                self.id2label[id_].append(label)
+
+    def _populate_menus(self):
+        id_ = self.store.ids[0]
+        if id_:
+            menu = create_dropdown_menu(self.store, list(self.id2label), "id")
+            menu.currentTextChanged.connect(self.refresh_label_menu)
+            self.menus["id"] = menu
+        self.menus["label"] = create_dropdown_menu(
+            self.store, self.id2label[id_], "label",
+        )
+
+    def _update_items(self):
+        id_ = self.store.ids[0]
+        if id_:
+            self.menus["id"].update_items(list(self.id2label))
+        self.menus["label"].update_items(self.id2label[id_])
+
     def _lock_current_keypoint(self):
         self._locked = not self._locked
         if self._locked:
@@ -833,5 +876,8 @@ def add_entry(self, name, color):
         self._layout.addWidget(
             LabelPair(color, name, self), alignment=Qt.AlignmentFlag.AlignLeft
         )
-        self._container.setLayout(self._layout)
-        self._container.update()
+
+    def reset(self):
+        self.scheme_dict = {}
+        for i in reversed(range(self._layout.count())):
+            self._layout.itemAt(i).widget().deleteLater()

src/napari_deeplabcut/keypoints.py

@@ -65,12 +65,21 @@ def default(cls):
 class KeypointStore:
     def __init__(self, viewer, layer: Points):
         self.viewer = viewer
+        self._keypoints = []
         self.layer = layer
+        self.viewer.dims.set_current_step(0, 0)
+
+    @property
+    def layer(self):
+        return self._layer
+
+    @layer.setter
+    def layer(self, layer):
+        self._layer = layer
         all_pairs = self.layer.metadata["header"].form_individual_bodypart_pairs()
         self._keypoints = [
             Keypoint(label, id_) for id_, label in all_pairs
         ]  # Ordered references to all possible keypoints
-        self.viewer.dims.set_current_step(0, 0)
 
     @property
     def current_step(self):