Interactive keypoint matching proofreading

MANIFEST.in

@@ -2,6 +2,10 @@ include LICENSE
 include README.md
 include src/napari_deeplabcut/styles/*.mplstyle
 include src/napari_deeplabcut/assets/napari_shortcuts.svg
+include src/napari_deeplabcut/assets/superanimal_topviewmouse.jpg
+include src/napari_deeplabcut/assets/superanimal_topviewmouse.json
+include src/napari_deeplabcut/assets/superanimal_quadruped.jpg
+include src/napari_deeplabcut/assets/superanimal_quadruped.json
 
 recursive-exclude * __pycache__
 recursive-exclude * *.py[co]

setup.cfg

@@ -40,6 +40,7 @@ install_requires =
     pyyaml
     qtpy>=2.4
     scikit-image
+    scipy
     tables
 python_requires = >=3.9
 include_package_data = True

src/napari_deeplabcut/_reader.py

@@ -1,4 +1,5 @@
 import glob
+import json
 import os
 from pathlib import Path
 from typing import Dict, List, Optional, Sequence
@@ -19,6 +20,10 @@
 SUPPORTED_VIDEOS = ".mp4", ".mov", ".avi"
 
 
+def is_video(filename: str):
+    return any(filename.lower().endswith(ext) for ext in SUPPORTED_VIDEOS)
+
+
 def get_hdf_reader(path):
     if isinstance(path, list):
         path = path[0]
@@ -152,6 +157,17 @@ def _populate_metadata(
     }
 
 
+def _load_superkeypoints_diagram(super_animal: str):
+    path = str(Path(__file__).parent / "assets" / f"{super_animal}.jpg")
+    return imread(path), {"root": ""}, "images"
+
+
+def _load_superkeypoints(super_animal: str):
+    path = str(Path(__file__).parent / "assets" / f"{super_animal}.json")
+    with open(path) as f:
+        return json.load(f)
+
+
 def _load_config(config_path: str):
     with open(config_path) as file:
         return yaml.safe_load(file)
@@ -171,6 +187,10 @@ def read_config(configname: str) -> List[LayerData]:
     metadata["ndim"] = 3
     metadata["property_choices"] = metadata.pop("properties")
     metadata["metadata"]["project"] = os.path.dirname(configname)
+    conversion_tables = config.get("SuperAnimalConversionTables")
+    if conversion_tables is not None:
+        super_animal, table = conversion_tables.popitem()
+        metadata["metadata"]["tables"] = {super_animal: table}
     return [(None, metadata, "points")]
 
 
@@ -299,7 +319,7 @@ def _read_frame(ind):
     elems[-1] = elems[-1].split(".")[0]
     root = os.path.join(*elems)
     params = {
-        "name": os.path.split(filename)[1],
+        "name": filename,
         "metadata": {
             "root": root,
         },

src/napari_deeplabcut/_widgets.py

@@ -31,7 +31,6 @@
     QCheckBox,
     QComboBox,
     QDialog,
-    QDialogButtonBox,
     QFileDialog,
     QGroupBox,
     QHBoxLayout,
@@ -49,7 +48,12 @@
 )
 
 from napari_deeplabcut import keypoints
-from napari_deeplabcut._reader import _load_config
+from napari_deeplabcut._reader import (
+    _load_config,
+    _load_superkeypoints_diagram,
+    _load_superkeypoints,
+    is_video,
+)
 from napari_deeplabcut._writer import _write_config, _write_image, _form_df
 from napari_deeplabcut.misc import (
     encode_categories,
@@ -349,18 +353,14 @@ def _update_buttons_checked(self) -> None:
                     QIcon(os.path.join(icon_dir, "Pan_checked.png"))
                 )
             else:
-                self._actions["pan"].setIcon(
-                    QIcon(os.path.join(icon_dir, "Pan.png"))
-                )
+                self._actions["pan"].setIcon(QIcon(os.path.join(icon_dir, "Pan.png")))
         if "zoom" in self._actions:
             if self._actions["zoom"].isChecked():
                 self._actions["zoom"].setIcon(
                     QIcon(os.path.join(icon_dir, "Zoom_checked.png"))
                 )
             else:
-                self._actions["zoom"].setIcon(
-                    QIcon(os.path.join(icon_dir, "Zoom.png"))
-                )
+                self._actions["zoom"].setIcon(QIcon(os.path.join(icon_dir, "Zoom.png")))
 
 
 class KeypointMatplotlibCanvas(QWidget):
@@ -602,6 +602,8 @@ def __init__(self, napari_viewer):
         self.video_widget.setVisible(False)
 
         # form helper display
+        self._keypoint_mapping_button = None
+        self._func_id = None
         help_buttons = self._form_help_buttons()
         self._layout.addLayout(help_buttons)
 
@@ -687,6 +689,72 @@ def __init__(self, napari_viewer):
     def settings(self):
         return QSettings()
 
+    def load_superkeypoints_diagram(self):
+        points_layer = None
+        for layer in self.viewer.layers:
+            if isinstance(layer, Points):
+                points_layer = layer
+                break
+
+        if points_layer is None:
+            return
+
+        tables = deepcopy(points_layer.metadata.get("tables", {}))
+        if not tables:
+            return
+
+        super_animal, table = tables.popitem()
+        layer_data = _load_superkeypoints_diagram(super_animal)
+        self.viewer.add_image(layer_data[0], metadata=layer_data[1])
+        superkpts_dict = _load_superkeypoints(super_animal)
+        xy = []
+        labels = []
+        for kpt_ref, kpt_super in table.items():
+            xy.append([0.0, *superkpts_dict[kpt_super]])
+            labels.append(kpt_ref)
+        points_layer.data = np.array(xy)
+        properties = deepcopy(points_layer.properties)
+        properties["label"] = np.array(labels)
+        points_layer.properties = properties
+        self._keypoint_mapping_button.setText("Map keypoints")
+        self._keypoint_mapping_button.clicked.disconnect(self._func_id)
+        self._keypoint_mapping_button.clicked.connect(
+            lambda: self._map_keypoints(super_animal)
+        )
+
+    def _map_keypoints(self, super_animal: str):
+        points_layer = None
+        for layer in self.viewer.layers:
+            if isinstance(layer, Points) and layer.metadata.get("tables"):
+                points_layer = layer
+                break
+
+        if points_layer is None or ~np.any(points_layer.data):
+            return
+
+        xy = points_layer.data[:, 1:3]
+        superkpts_dict = _load_superkeypoints(super_animal)
+        xy_ref = np.c_[[val for val in superkpts_dict.values()]]
+        neighbors = keypoints._find_nearest_neighbors(xy, xy_ref)
+        found = neighbors != -1
+        if ~np.any(found):
+            return
+
+        project_path = points_layer.metadata["project"]
+        config_path = str(Path(project_path) / "config.yaml")
+        cfg = _load_config(config_path)
+        conversion_tables = cfg.get("SuperAnimalConversionTables", {})
+        conversion_tables[super_animal] = dict(
+            zip(
+                map(
+                    str, points_layer.metadata["header"].bodyparts
+                ),  # Needed to fix an ugly yaml RepresenterError
+                map(str, list(np.array(list(superkpts_dict))[neighbors[found]])),
+            )
+        )
+        _write_config(config_path, cfg)
+        self.viewer.status = "Mapping to superkeypoint set successfully saved"
+
     def start_tutorial(self):
         Tutorial(self.viewer.window._qt_window.current()).show()
 
@@ -753,6 +821,12 @@ def _form_help_buttons(self):
         tutorial = QPushButton("Start tutorial")
         tutorial.clicked.connect(self.start_tutorial)
         layout.addWidget(tutorial)
+        self._keypoint_mapping_button = QPushButton("Load superkeypoints diagram")
+        self._func_id = self._keypoint_mapping_button.clicked.connect(
+            self.load_superkeypoints_diagram
+        )
+        self._keypoint_mapping_button.hide()
+        layout.addWidget(self._keypoint_mapping_button)
         return layout
 
     def _extract_single_frame(self, *args):
@@ -780,7 +854,7 @@ def _extract_single_frame(self, *args):
                         "properties": points_layer.properties,
                     },
                 )
-                df = df.iloc[ind: ind + 1]
+                df = df.iloc[ind : ind + 1]
                 df.index = pd.MultiIndex.from_tuples([Path(output_path).parts[-3:]])
                 filepath = os.path.join(
                     image_layer.metadata["root"], "machinelabels-iter0.h5"
@@ -889,7 +963,9 @@ def to_hex(nparray):
                 self._display.update_color_scheme(
                     {
                         name: to_hex(color)
-                        for name, color in layer.metadata["face_color_cycles"][mode].items()
+                        for name, color in layer.metadata["face_color_cycles"][
+                            mode
+                        ].items()
                     }
                 )
 
@@ -933,7 +1009,7 @@ def on_insert(self, event):
         logging.debug(f"Inserting Layer {layer}")
         if isinstance(layer, Image):
             paths = layer.metadata.get("paths")
-            if paths is None:  # Then it's a video file
+            if paths is None and is_video(layer.name):
                 self.video_widget.setVisible(True)
             # Store the metadata and pass them on to the other layers
             self._images_meta.update(
@@ -986,15 +1062,16 @@ def on_insert(self, event):
                         "face_color_cycles"
                     ]
                     _layer.face_color = "label"
-                    _layer.face_color_cycle = new_metadata["face_color_cycles"][
-                        "label"
-                    ]
+                    _layer.face_color_cycle = new_metadata["face_color_cycles"]["label"]
                     _layer.events.face_color()
                     store.layer = _layer
                 self._update_color_scheme()
 
                 return
 
+            if layer.metadata.get("tables", ""):
+                self._keypoint_mapping_button.show()
+
             store = keypoints.KeypointStore(self.viewer, layer)
             self._stores[layer] = store
             # TODO Set default dir of the save file dialog
@@ -1073,9 +1150,9 @@ def on_remove(self, event):
 
     def on_active_layer_change(self, event) -> None:
         """Updates the GUI when the active layer changes
-            * Hides all KeypointsDropdownMenu that aren't for the selected layer
-            * Sets the visibility of the "Color mode" box to True if the selected layer
-                is a multi-animal one, or False otherwise
+        * Hides all KeypointsDropdownMenu that aren't for the selected layer
+        * Sets the visibility of the "Color mode" box to True if the selected layer
+            is a multi-animal one, or False otherwise
         """
         self._color_grp.setVisible(self._is_multianimal(event.value))
         menu_idx = -1
@@ -1092,7 +1169,8 @@ def _update_colormap(self, colormap_name):
         for layer in self.viewer.layers.selection:
             if isinstance(layer, Points) and layer.metadata:
                 face_color_cycle_maps = build_color_cycles(
-                    layer.metadata["header"], colormap_name,
+                    layer.metadata["header"],
+                    colormap_name,
                 )
                 layer.metadata["face_color_cycles"] = face_color_cycle_maps
                 face_color_prop = "label"
@@ -1110,9 +1188,8 @@ def cycle_through_label_modes(self, *args):
 
     @register_points_action("Change color mode")
     def cycle_through_color_modes(self, *args):
-        if (
-            self._active_layer_is_multianimal()
-            or self.color_mode != str(keypoints.ColorMode.BODYPART)
+        if self._active_layer_is_multianimal() or self.color_mode != str(
+            keypoints.ColorMode.BODYPART
         ):
             self.color_mode = next(keypoints.ColorMode)
 
@@ -1152,7 +1229,8 @@ def color_mode(self, mode: Union[str, keypoints.ColorMode]):
             if isinstance(layer, Points) and layer.metadata:
                 layer.face_color = face_color_mode
                 layer.face_color_cycle = layer.metadata["face_color_cycles"][
-                    face_color_mode]
+                    face_color_mode
+                ]
                 layer.events.face_color()
 
         for btn in self._color_mode_selector.buttons():