Rewrite logic of keypoint menus

Author: jeylau

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()
 
@@ -349,29 +349,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()
@@ -434,6 +428,22 @@ def on_insert(self, event):
                 10, partial(self._move_image_layer_to_bottom, event.index)
             )
         elif isinstance(layer, Points):
+            # If the Points layer comes from a config file and some Points layers
+            # were already added, then we only update existing store's metadata.
+            if layer.metadata.get("project", "") and self._stores:
+                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()
+                return
+
             store = keypoints.KeypointStore(self.viewer, layer)
             self._stores[layer] = store
             # TODO Set default dir of the save file dialog
@@ -469,7 +479,8 @@ 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._stores.pop(layer, None)
@@ -528,7 +539,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)
@@ -538,6 +549,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__(
@@ -549,22 +564,11 @@ def __init__(
         self.store = store
         self.store.layer.events.current_properties.connect(self.update_menus)
 
-        # 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")
@@ -575,6 +579,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 update_menus(self, event):
         keypoint = self.store.current_keypoint
         for attr, menu in self.menus.items():
@@ -805,5 +832,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

@@ -59,12 +59,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):