Ensure Annotations layer exists & bring prediction layer creation int…

…o helper function
fractal-napari-plugins-collection · Sep 16, 2024 · cd7e017 · cd7e017
1 parent e778a75
commit cd7e017
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 23 deletions.
diff --git a/src/napari_feature_classifier/annotator_widget.py b/src/napari_feature_classifier/annotator_widget.py
@@ -136,16 +136,7 @@ def __init__(
         for layer in self._viewer.layers:
             if type(layer) == napari.layers.Labels and layer.name == "Annotations":
                 self._viewer.layers.remove(layer)
-        self._annotations_layer = self._viewer.add_labels(
-            self._last_selected_label_layer.data,
-            scale=self._last_selected_label_layer.scale,
-            name="Annotations",
-            translate=self._last_selected_label_layer.translate,
-        )
-        self._annotations_layer.editable = False
-
-        # Set the label selection to a valid label layer => Running into proxy bug
-        self._viewer.layers.selection.active = self._last_selected_label_layer
+        self.add_annotations_layer()
 
         # Class selection
         self.ClassSelection = ClassSelection  # pylint: disable=C0103
@@ -212,11 +203,26 @@ def selection_changed(self, event):
             self._save_destination.enabled = False
             self._class_selector.enabled = False
 
+    def add_annotations_layer(self):
+        self._annotations_layer = self._viewer.add_labels(
+            self._last_selected_label_layer.data,
+            scale=self._last_selected_label_layer.scale,
+            name="Annotations",
+            translate=self._last_selected_label_layer.translate,
+        )
+        self._annotations_layer.editable = False
+        # Set the label selection to a valid label layer
+        self._viewer.layers.selection.active = self._last_selected_label_layer
+
     def toggle_label(self, labels_layer, event):
         """
         Callback for when a label is clicked. It then updates the color of that
         label in the annotation layer.
         """
+        # If the annotations layer is missing, add it back
+        if "Annotations" not in [x.name for x in self._viewer.layers]:
+            self.add_annotations_layer()
+
         # Need to translate & scale position that event.position returns by the
         # label_layer scale.
         # If scale is (1, 1, 1), nothing changes

diff --git a/src/napari_feature_classifier/classifier_widget.py b/src/napari_feature_classifier/classifier_widget.py
@@ -251,13 +251,7 @@ def __init__(
         for layer in self._viewer.layers:
             if type(layer) == napari.layers.Labels and layer.name == "Predictions":
                 self._viewer.layers.remove(layer)
-        self._prediction_layer = self._viewer.add_labels(
-            self._last_selected_label_layer.data,
-            scale=self._last_selected_label_layer.scale,
-            name="Predictions",
-            translate=self._last_selected_label_layer.translate,
-        )
-        self._prediction_layer.contour = 2
+        self.add_prediction_layer()
 
         # Set the label selection to a valid label layer => Running into proxy bug
         self._viewer.layers.selection.active = self._last_selected_label_layer
@@ -343,6 +337,15 @@ def add_features_to_classifier(self):
                     dict_of_features[layer.name] = layer.features
         self._classifier.add_dict_of_features(dict_of_features)
 
+    def add_prediction_layer(self):
+        self._prediction_layer = self._viewer.add_labels(
+            self._last_selected_label_layer.data,
+            scale=self._last_selected_label_layer.scale,
+            name="Predictions",
+            translate=self._last_selected_label_layer.translate,
+        )
+        self._prediction_layer.contour = 2
+
     def make_predictions(self):
         """
         Make predictions for all relevant label layers and add them to the
@@ -446,12 +449,7 @@ def _init_prediction_layer(
             "Predictions" not in [x.name for x in self._viewer.layers]
             and ensure_layer_presence
         ):
-            self._prediction_layer = self._viewer.add_labels(
-                self._last_selected_label_layer.data,
-                scale=self._last_selected_label_layer.scale,
-                name="Predictions",
-                translate=self._last_selected_label_layer.translate,
-            )
+            self.add_prediction_layer()
         if ensure_layer_presence:
             # Ensure correct layer order: This sometimes fails with weird
             # EmitLoopError & IndexError that should be ignored