Add colormap selector

Author: jeylau

src/napari_deeplabcut/_reader.py

@@ -124,12 +124,7 @@ def _populate_metadata(
         ids = header.individuals
     if likelihood is None:
         likelihood = np.ones(len(labels))
-    label_colors = misc.build_color_cycle(len(header.bodyparts), colormap)
-    id_colors = misc.build_color_cycle(len(header.individuals), colormap)
-    face_color_cycle_maps = {
-        "label": dict(zip(header.bodyparts, label_colors)),
-        "id": dict(zip(header.individuals, id_colors)),
-    }
+    face_color_cycle_maps = misc.build_color_cycles(header, colormap)
     face_color_prop = "id" if ids[0] else "label"
     return {
         "name": "keypoints",
@@ -151,6 +146,7 @@ def _populate_metadata(
         "metadata": {
             "header": header,
             "face_color_cycles": face_color_cycle_maps,
+            "colormap_name": colormap,
             "paths": paths or [],
         },
     }

src/napari_deeplabcut/_widgets.py

@@ -4,6 +4,7 @@
 from datetime import datetime
 from functools import partial, cached_property
 from math import ceil, log10
+import matplotlib.pyplot as plt
 import pandas as pd
 from pathlib import Path
 from types import MethodType
@@ -49,6 +50,7 @@
     encode_categories,
     to_os_dir_sep,
     guarantee_multiindex_rows,
+    build_color_cycles
 )
 
 
@@ -701,6 +703,11 @@ def on_insert(self, event):
             # Hide out of slice checkbox
             point_controls.outOfSliceCheckBox.hide()
             point_controls.layout().itemAt(15).widget().hide()
+            # Add dropdown menu for colormap picking
+            colormap_selector = DropdownMenu(plt.colormaps, self)
+            colormap_selector.update_to(layer.metadata["colormap_name"])
+            colormap_selector.currentTextChanged.connect(self._update_colormap)
+            point_controls.layout().addRow("colormap", colormap_selector)
 
         for layer_ in self.viewer.layers:
             if not isinstance(layer_, Image):
@@ -732,6 +739,20 @@ def on_remove(self, event):
             self._trail_cb.setChecked(False)
             self._trails = None
 
+    def _update_colormap(self, colormap_name):
+        for layer in self.viewer.layers:
+            if isinstance(layer, Points) and layer.metadata:
+                face_color_cycle_maps = build_color_cycles(
+                    layer.metadata["header"], colormap_name,
+                )
+                layer.metadata["face_color_cycles"] = face_color_cycle_maps
+                face_color_prop = layer._face.color_properties.name
+                layer.face_color = face_color_prop
+                layer.face_color_cycle = face_color_cycle_maps[face_color_prop]
+                layer.events.face_color()
+                self._update_color_scheme()
+                break
+
     @register_points_action("Change labeling mode")
     def cycle_through_label_modes(self, *args):
         self.label_mode = next(keypoints.LabelMode)
@@ -1127,4 +1148,5 @@ def add_entry(self, name, color):
     def reset(self):
         self.scheme_dict = {}
         for i in reversed(range(self._layout.count())):
-            self._layout.itemAt(i).widget().deleteLater()
+            w = self._layout.itemAt(i).widget()
+            self._layout.removeWidget(w)

src/napari_deeplabcut/misc.py

@@ -101,6 +101,15 @@ def build_color_cycle(n_colors: int, colormap: Optional[str] = "viridis") -> np.
     return cmap.map(np.linspace(0, 1, n_colors))
 
 
+def build_color_cycles(header: DLCHeader, colormap: Optional[str] = "viridis"):
+    label_colors = build_color_cycle(len(header.bodyparts), colormap)
+    id_colors = build_color_cycle(len(header.individuals), colormap)
+    return {
+        "label": dict(zip(header.bodyparts, label_colors)),
+        "id": dict(zip(header.individuals, id_colors)),
+    }
+
+
 class DLCHeader:
     def __init__(self, columns: pd.MultiIndex):
         self.columns = columns