diff --git a/src/napari_deeplabcut/_reader.py b/src/napari_deeplabcut/_reader.py index 5a9754b..ca06849 100644 --- a/src/napari_deeplabcut/_reader.py +++ b/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 diff --git a/src/napari_deeplabcut/_widgets.py b/src/napari_deeplabcut/_widgets.py index 8111067..608378c 100644 --- a/src/napari_deeplabcut/_widgets.py +++ b/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")