Seb/display skeleton

src/napari_deeplabcut/_reader.py

@@ -227,7 +227,34 @@ def read_hdf(filename: str) -> List[LayerData]:
         metadata["metadata"]["root"] = os.path.split(filename)[0]
         # Store file name in case the layer's name is edited by the user
         metadata["metadata"]["name"] = metadata["name"]
+
+        limb_pairs = [['nose','LeftForelimb'], 
+                      ['nose', 'RightForelimb'], 
+                      ['LeftForelimb', 'RightForelimb'], 
+                      ['LeftHindlimb', 'RightHindlimb'], 
+                      ['TailBase', 'LeftHindlimb'], 
+                      ['TailBase', 'RightHindlimb'],
+                      ['TailBase', 'Tail1'],
+                      ['Tail1', 'Tail2'],
+                      ['Tail2', 'Tail3']
+        ]
+        n = temp.shape[0]
+        vectors = np.zeros((n*len(limb_pairs), 2, 3))
+        for i, (kpt1, kpt2) in enumerate(limb_pairs):
+            origin = temp.xs(kpt1, level='bodyparts', axis=1).to_numpy()[:,:2]
+            end = temp.xs(kpt2, level='bodyparts', axis=1).to_numpy()[:,:2]
+            vec = end-origin
+            vectors[i*n:(i+1)*n, 0, [2,1]] = origin
+            vectors[i*n:(i+1)*n, 1, [2,1]] = vec
+            vectors[i*n:(i+1)*n, :, 0] = np.arange(temp.shape[0])[:, None]
+
         layers.append((data, metadata, "points"))
+        # Make a dictionary of the colors of the bodyparts based on the colormap
+        # and the bodyparts
+
+        layers.append((vectors, {'edge_width':1, 'edge_color':'yellow'}, "vectors"))
+
+
     return layers
 
 

src/napari_deeplabcut/_widgets.py

@@ -290,7 +290,6 @@ def on_close(self, event, widget):
     else:
         event.accept()
 
-
 class KeypointControls(QWidget):
     def __init__(self, napari_viewer):
         super().__init__()
@@ -354,10 +353,21 @@ def __init__(self, napari_viewer):
         self._trail_cb.stateChanged.connect(self._show_trails)
         self._trails = None
 
+        # Add checkbox to show skeleton
+        skeleton_label = QLabel("Show skeleton")
+        self._skeleton_cb = QCheckBox()
+        self._skeleton_cb.setToolTip("toggle skeleton visibility")
+        self._skeleton_cb.setChecked(False)
+        self._skeleton_cb.setEnabled(False)
+        self._skeleton_cb.stateChanged.connect(self._show_skeleton)
+        self._skeleton = None
+
         self._view_scheme_cb = QCheckBox("Show color scheme", parent=self)
 
-        hlayout.addWidget(trail_label)
+        hlayout.addWidget(self._skeleton_cb)
+        hlayout.addWidget(skeleton_label)
         hlayout.addWidget(self._trail_cb)
+        hlayout.addWidget(trail_label)
         hlayout.addWidget(self._view_scheme_cb)
 
         self._layout.addLayout(hlayout)
@@ -394,6 +404,8 @@ def __init__(self, napari_viewer):
             QTimer.singleShot(10, self.start_tutorial)
             self.settings.setValue("first_launch", False)
 
+
+
     @cached_property
     def settings(self):
         return QSettings()
@@ -427,6 +439,13 @@ def _show_trails(self, state):
             self._trails.visible = True
         elif self._trails is not None:
             self._trails.visible = False
+
+
+    def _show_skeleton(self, state):
+        if state == Qt.Checked:
+            # Check if "skeleton" and "skeleton_color" are in the config.yaml metadata
+            return True
+
 
     def _form_video_action_menu(self):
         group_box = QGroupBox("Video")
@@ -681,6 +700,7 @@ def on_insert(self, event):
                 }
             )
             self._trail_cb.setEnabled(True)
+            self._skeleton_cb.setEnabled(True)
 
             # Hide the color pickers, as colormaps are strictly defined by users
             controls = self.viewer.window.qt_viewer.dockLayerControls
@@ -710,6 +730,7 @@ def on_remove(self, event):
                 menu.deleteLater()
                 menu.destroy()
             self._trail_cb.setEnabled(False)
+            self._skeleton_cb.setEnabled(False)
             self.last_saved_label.hide()
         elif isinstance(layer, Image):
             self._images_meta = dict()
@@ -718,6 +739,7 @@ def on_remove(self, event):
                 self.video_widget.setVisible(False)
         elif isinstance(layer, Tracks):
             self._trail_cb.setChecked(False)
+            self._skeleton_cb.setChecked(False)
             self._trails = None
 
     @register_points_action("Change labeling mode")