diff --git a/docs/source/user_guide/gui.md b/docs/source/user_guide/gui.md index 022ad3f821..3cb12f069a 100644 --- a/docs/source/user_guide/gui.md +++ b/docs/source/user_guide/gui.md @@ -203,6 +203,22 @@ And for a bounding boxes dataset, you will see a view more like the one below: Note the additional bounding boxes layer that is loaded for bounding boxes datasets. For both poses and bounding boxes datasets, you can toggle the visibility of any of these layers by clicking on the eye icon. +### Reset pose edits + +In **Edit tracked data**, use **Reset current frame** to restore dragged or +deleted pose points in the displayed frame. Edits in other frames are kept. +Use **Reset all frames** to restore all pose points in the selected dataset. + +Both controls restore the dataset as it was loaded into the viewer. Any edits +already saved in the loaded file are retained, including their edited flags. +Resetting does not change the file on disk. Save the dataset separately to +persist the restored state. + +Reset is available for pose datasets with their companion tracks layer still +present. The current frame control is disabled when the time axis is displayed +as a spatial axis. Bounding boxes and individual edit history are not supported. +The loaded baseline is held in memory for each pose layer. + ### The points layer The points layer shows the data for the current frame. diff --git a/movement/napari/edit_timeline_widget.py b/movement/napari/edit_timeline_widget.py index f59748d242..6f1ed6eb0c 100644 --- a/movement/napari/edit_timeline_widget.py +++ b/movement/napari/edit_timeline_widget.py @@ -16,13 +16,24 @@ from napari.utils.theme import get_theme from napari.viewer import Viewer from qtpy.QtCore import QTimer, Signal -from qtpy.QtWidgets import QCheckBox, QLabel, QVBoxLayout, QWidget +from qtpy.QtWidgets import ( + QCheckBox, + QLabel, + QMessageBox, + QPushButton, + QVBoxLayout, + QWidget, +) from movement.napari.layer_wiring import ( MAX_FRAME_IDX_KEY, + POINTS_POSITION_KEY, POINTS_PROPERTIES_KEY, + TRACKS_LAYER_KEY, active_movement_points_layer, + frame_axis_is_sliced, is_movement_points_layer, + reset_edits, ) if TYPE_CHECKING: @@ -51,9 +62,10 @@ class EditControlsWidget(QWidget): show_individuals_toggled = Signal(bool) - def __init__(self, parent=None): + def __init__(self, parent=None, napari_viewer=None): """Initialise the instructions label and display checkbox.""" super().__init__(parent=parent) + self.viewer = napari_viewer instructions = QLabel( "Use the points layer controls to move or delete keypoints. " "Frames with edited points are flagged as coloured bars on the " @@ -73,8 +85,75 @@ def __init__(self, parent=None): layout = QVBoxLayout() layout.addWidget(instructions) layout.addWidget(self.show_individuals_checkbox) + self.reset_frame_button = QPushButton("Reset current frame") + self.reset_all_button = QPushButton("Reset all frames") + for button in (self.reset_frame_button, self.reset_all_button): + button.setToolTip( + "Restore pose points to the file as loaded, including " + "any previously saved edits." + ) + layout.addWidget(button) + self.reset_frame_button.clicked.connect(self._reset_frame) + self.reset_all_button.clicked.connect(self._reset_all) + if self.viewer is not None: + for event in ( + self.viewer.layers.events.inserted, + self.viewer.layers.events.removed, + self.viewer.layers.selection.events.active, + self.viewer.dims.events.order, + self.viewer.dims.events.ndisplay, + ): + event.connect(self._schedule_reset_enabled) + self._update_reset_enabled() self.setLayout(layout) + def _reset_layer(self): + """Find a loaded pose layer with its companion still in the viewer.""" + if self.viewer is None: + return None + layer = active_movement_points_layer(self.viewer) + if ( + layer is not None + and POINTS_POSITION_KEY in layer.metadata + and layer.metadata.get(TRACKS_LAYER_KEY) in self.viewer.layers + ): + return layer + return None + + def _schedule_reset_enabled(self, event=None): + """Wait until loading both layers has finished.""" + QTimer.singleShot(0, self._update_reset_enabled) + + def _update_reset_enabled(self): + """Disable restoration when no supported pose layer is loaded.""" + enabled = self._reset_layer() is not None + self.reset_all_button.setEnabled(enabled) + self.reset_frame_button.setEnabled( + enabled and frame_axis_is_sliced(self.viewer) + ) + + def _reset_frame(self): + """Reset only the current frame, leaving other frames unchanged.""" + layer = self._reset_layer() + if layer is not None and frame_axis_is_sliced(self.viewer): + reset_edits(layer, self.viewer.dims.current_step[0]) + + def _reset_all(self): + """Reset all frames in the active pose layer.""" + layer = self._reset_layer() + if layer is not None: + answer = QMessageBox.warning( + self, + title="Reset all pose edits?", + text="Discard all corrections made since this dataset was " + "loaded? Previously saved edits will be kept. " + "This reset cannot be undone.", + buttons=QMessageBox.Reset | QMessageBox.Cancel, + defaultButton=QMessageBox.Cancel, + ) + if answer == QMessageBox.Reset: + reset_edits(layer) + class EditTimelineWidget(QWidget): """Dock widget flagging frames with edited points. @@ -283,6 +362,21 @@ def _on_layer_data_changed(self, event): """ if event.source is not self.active_layer: return + if getattr(event, "movement_reset", False) is True: + frame = event.frame + previous = self._reconstruct_previously_removed_points( + self.active_layer + ) + if frame is None: + self._removed_points = previous + else: + self._removed_points = [ + point + for point in self._removed_points + if point[0] != frame + ] + [point for point in previous if point[0] == frame] + self._redraw_bars() + return if event.action == ActionType.REMOVING: self._capture_removed_points(event) self._redraw_bars() diff --git a/movement/napari/layer_wiring.py b/movement/napari/layer_wiring.py index c79027192e..8f27029bc4 100644 --- a/movement/napari/layer_wiring.py +++ b/movement/napari/layer_wiring.py @@ -18,23 +18,29 @@ from weakref import WeakSet import numpy as np +import pandas as pd from napari.components.dims import RangeTuple from napari.layers import Points from napari.layers.base import ActionType -from movement.napari.layer_styles import EDITED_POINT_SYMBOL +from movement.napari.layer_styles import EDITED_POINT_SYMBOL, PointsStyle # Metadata keys stored on the movement Points layer. # - POINTS_LAYER_KEY marks the layer as movement-created. # - POINTS_PROPERTIES_KEY holds the full properties df, incl. the NaN rows # dropped from the live layer, needed to reconstruct the dataset. +# - POINTS_POSITION_KEY holds the full Tracks array as loaded, incl. the +# NaN rows, row-aligned with POINTS_PROPERTIES_KEY. Together they are the +# loaded state that edits can be reset to. Poses datasets only. # - DATASET_ATTRS_KEY holds the source dataset's attrs (source_software, fps…). # - TRACKS_LAYER_KEY holds a reference to the companion Tracks layer. # - MAX_FRAME_IDX_KEY holds the last frame index of the source data, # including leading/trailing all-NaN frames (which are dropped from the # napari layer data array). +# The loaded-state values are never modified after loading. POINTS_LAYER_KEY: str = "movement_points_layer" POINTS_PROPERTIES_KEY: str = "movement_points_properties" +POINTS_POSITION_KEY: str = "movement_points_position" DATASET_ATTRS_KEY: str = "movement_dataset_attrs" TRACKS_LAYER_KEY: str = "movement_tracks_layer" MAX_FRAME_IDX_KEY: str = "movement_max_frame_idx" @@ -186,6 +192,8 @@ def on_points_data_changed(event): layer = event.source if not isinstance(layer, Points): return + if getattr(event, "movement_reset", False) is True: + return if event.action == ActionType.CHANGED: moved_indices = list(event.data_indices) @@ -267,3 +275,90 @@ def set_tracks_layer_data(tracks_layer, data, properties): tracks_layer.data = data tracks_layer.properties = properties tracks_layer.color_by = color_by + + +def points_layer_properties(properties: pd.DataFrame) -> pd.DataFrame: + """Drop the internal columns not needed in Points layer properties. + + The ``_factorized`` columns are used for Tracks/Shapes colouring and + ``position_is_nan`` for reconstructing the dataset on save; neither + is needed in the Points layer tooltips. + """ + return properties.loc[ + :, + ~properties.columns.str.endswith("_factorized") + & (properties.columns != "position_is_nan"), + ] + + +def reset_edits(points_layer: Points, frame: int | None = None) -> None: + """Restore loaded pose points for one frame, or all frames if None. + + The loaded state is read from ``POINTS_POSITION_KEY`` (positions) + and ``POINTS_PROPERTIES_KEY`` (confidence, ``edited`` flags and + which points were present), so it includes edits saved in the input + file. Deleted points are restored too. Other frames retain their + current data. Identity swaps and individual edit history are not + supported. + """ + loaded_data = points_layer.metadata[POINTS_POSITION_KEY] + loaded_props = points_layer.metadata[POINTS_PROPERTIES_KEY] + tracks = points_layer.metadata[TRACKS_LAYER_KEY] + # NaN rows are never shown in the layers, so never restore them + restore = ~loaded_props["position_is_nan"].to_numpy() + keep = np.zeros(len(tracks.data), dtype=bool) + if frame is not None: + restore &= loaded_data[:, 1] == frame + keep = tracks.data[:, 1] != frame + data = np.concatenate([tracks.data[keep], loaded_data[restore]]) + order = np.lexsort((data[:, 1], data[:, 0])) + data = data[order] + + def merge_properties(current, loaded: pd.DataFrame): + keys = current.keys() | set(loaded.columns) + return { + key: np.concatenate( + [ + np.asarray( + current.get(key, np.zeros(len(tracks.data), bool)) + )[keep], + loaded[key].to_numpy()[restore] + if key in loaded + else np.zeros(restore.sum(), bool), + ] + )[order] + for key in keys + } + + properties = merge_properties( + points_layer.properties, points_layer_properties(loaded_props) + ) + track_properties = merge_properties(tracks.properties, loaded_props) + # Kept points keep their current symbol; restored points are derived + # from their loaded edited flag, as on load + loaded_edited = ( + loaded_props["edited"].to_numpy()[restore] + if "edited" in loaded_props + else np.zeros(restore.sum(), bool) + ) + symbols = np.concatenate( + [ + np.asarray(points_layer.symbol, dtype=object)[keep], + np.where(loaded_edited, EDITED_POINT_SYMBOL, PointsStyle.symbol), + ] + )[order] + # Avoid interpreting restoration as a user drag or deletion. Notify + # observers once, after Points and Tracks are consistent again. + with points_layer.events.data.blocker(): + points_layer.data = data[:, 1:].copy() + points_layer.properties = properties + points_layer.symbol = symbols + points_layer.selected_data = set() + set_tracks_layer_data(tracks, data, track_properties) + points_layer.events.data( + value=points_layer.data, + action=ActionType.CHANGED, + data_indices=(), + movement_reset=True, + frame=frame, + ) diff --git a/movement/napari/loader_widgets.py b/movement/napari/loader_widgets.py index 90c7d0da7f..ac6d714afa 100644 --- a/movement/napari/loader_widgets.py +++ b/movement/napari/loader_widgets.py @@ -28,11 +28,13 @@ DATASET_ATTRS_KEY, MAX_FRAME_IDX_KEY, POINTS_LAYER_KEY, + POINTS_POSITION_KEY, POINTS_PROPERTIES_KEY, TRACKS_LAYER_KEY, connect_viewer_callbacks, frame_axis_is_sliced, on_points_data_changed, + points_layer_properties, set_point_symbol_by_edited, update_frame_slider_range, ) @@ -339,23 +341,20 @@ def _add_points_layer(self): properties_df=self.properties, ) - # Filter out columns used internally (for Tracks/Shapes coloring, - # or for reconstructing the dataset on save) but not needed in - # Points layer tooltips: _factorized columns and position_is_nan. - points_properties = self.properties.loc[ - :, - ~self.properties.columns.str.endswith("_factorized") - & (self.properties.columns != "position_is_nan"), - ] + points_properties = points_layer_properties(self.properties) + metadata = { + MAX_FRAME_IDX_KEY: max(self.data[:, 1]), + POINTS_LAYER_KEY: True, + POINTS_PROPERTIES_KEY: self.properties, + DATASET_ATTRS_KEY: self.ds_attrs, + } + # Keep the loaded positions so edits can be reset (poses only) + if self.data_bboxes is None: + metadata[POINTS_POSITION_KEY] = self.data self.points_layer = self.viewer.add_points( self.data[self.data_not_nan, 1:], properties=points_properties.iloc[self.data_not_nan, :], - metadata={ - MAX_FRAME_IDX_KEY: max(self.data[:, 1]), - POINTS_LAYER_KEY: True, - POINTS_PROPERTIES_KEY: self.properties, - DATASET_ATTRS_KEY: self.ds_attrs, - }, + metadata=metadata, **points_style.as_kwargs(), ) self.points_layer.events.data.connect(on_points_data_changed) diff --git a/movement/napari/meta_widget.py b/movement/napari/meta_widget.py index f2ec074001..fcbfdb4366 100644 --- a/movement/napari/meta_widget.py +++ b/movement/napari/meta_widget.py @@ -47,7 +47,9 @@ def __init__(self, napari_viewer: Viewer, parent=None): # A collapsible "edit controls" widget that can be used # to show/hide and configure the edit timeline docked # to the bottom of the viewer. - self.edit_controls = EditControlsWidget(parent=self) + self.edit_controls = EditControlsWidget( + parent=self, napari_viewer=napari_viewer + ) self.edit_controls.show_individuals_toggled.connect( self._on_show_individuals_toggled ) diff --git a/tests/test_unit/test_napari_plugin/test_reset_edits.py b/tests/test_unit/test_napari_plugin/test_reset_edits.py new file mode 100644 index 0000000000..2de844d371 --- /dev/null +++ b/tests/test_unit/test_napari_plugin/test_reset_edits.py @@ -0,0 +1,245 @@ +"""Reset dragged and deleted pose points to their loaded state.""" + +import numpy as np +import pandas as pd +import pytest +from napari.components import ViewerModel +from qtpy.QtWidgets import QMessageBox + +from movement.napari import layer_wiring as wiring + + +@pytest.fixture +def reset_layers(): + """Create two tracks with two frames and their loaded state.""" + viewer = ViewerModel() + data = np.array( + [[0, 0, 1, 2], [0, 1, 3, 4], [1, 0, 5, 6], [1, 1, 7, 8]], dtype=float + ) + props = { + "confidence": np.array([0.1, 0.2, 0.3, 0.4]), + "individual": np.array(["a", "a", "b", "b"]), + "edited": np.array([True, False, False, False]), + } + loaded_props = pd.DataFrame(props).assign(position_is_nan=False) + points = viewer.add_points( + data[:, 1:].copy(), + properties=props, + metadata={ + wiring.POINTS_LAYER_KEY: True, + wiring.POINTS_PROPERTIES_KEY: loaded_props, + wiring.POINTS_POSITION_KEY: data.copy(), + }, + ) + tracks = viewer.add_tracks(data.copy(), properties=props) + points.metadata[wiring.TRACKS_LAYER_KEY] = tracks + wiring.set_point_symbol_by_edited(points) + points.events.data.connect(wiring.on_points_data_changed) + return points, tracks, data, props, viewer + + +@pytest.mark.parametrize("frame", [None, 0, 1]) +def test_reset_restores_deleted_points_and_keeps_other_edits( + reset_layers, frame +): + """Restore the selected frame without undoing another frame's edits.""" + points, tracks, original, props, viewer = reset_layers + points.data[1, 1:] = [100, 200] + points.events.data(action="changed", data_indices=(1,)) + points.remove([0, 2]) + wiring.reset_edits(points, frame=frame) + if frame in (None, 0): + assert len(points.data) == 4 + np.testing.assert_array_equal(tracks.data[[0, 2]], original[[0, 2]]) + else: + assert len(points.data) == 2 + if frame in (None, 1): + np.testing.assert_array_equal( + points.data[points.data[:, 0] == 1], original[[1, 3], 1:] + ) + else: + assert points.properties["edited"][1] + np.testing.assert_array_equal(points.data[1, 1:], [100, 200]) + expected = points.data.copy() + wiring.reset_edits(points, frame=frame) + np.testing.assert_array_equal(points.data, expected) + + +@pytest.mark.parametrize("frame", [None, 2]) +def test_reset_skips_missing_loaded_positions(reset_layers, frame): + """An originally missing point must stay absent after reset.""" + points, tracks, original, props, viewer = reset_layers + loaded_props = points.metadata[wiring.POINTS_PROPERTIES_KEY] + points.metadata[wiring.POINTS_PROPERTIES_KEY] = pd.concat( + [loaded_props, loaded_props.iloc[[-1]].assign(position_is_nan=True)], + ignore_index=True, + ) + loaded_positions = np.vstack([original, [1, 2, np.nan, np.nan]]) + points.metadata[wiring.POINTS_POSITION_KEY] = loaded_positions.copy() + wiring.reset_edits(points, frame=frame) + np.testing.assert_array_equal(tracks.data, original) + np.testing.assert_array_equal( + points.metadata[wiring.POINTS_POSITION_KEY], loaded_positions + ) + assert "position_is_nan" not in points.properties + + +def test_reset_all_restores_loaded_properties_and_symbols(reset_layers): + """Keep edits that were already in the loaded file, including styling.""" + points, tracks, original, props, viewer = reset_layers + symbols = points.symbol.copy() + points.remove([0, 1, 2]) + wiring.reset_edits(points) + np.testing.assert_array_equal(tracks.data, original) + for key, values in props.items(): + np.testing.assert_array_equal(points.properties[key], values) + np.testing.assert_array_equal(points.symbol, symbols) + + +def test_reset_controls_and_timeline(reset_layers, qtbot, monkeypatch): + """Buttons restore points and refresh timeline flags without OpenGL.""" + from movement.napari.edit_timeline_widget import ( + EditControlsWidget, + EditTimelineWidget, + ) + + points, tracks, original, props, viewer = reset_layers + monkeypatch.setattr( + QMessageBox, "warning", lambda *a, **kw: QMessageBox.Reset + ) + controls = EditControlsWidget(napari_viewer=viewer) + timeline = EditTimelineWidget(viewer) + qtbot.addWidget(controls) + qtbot.addWidget(timeline) + assert controls.reset_frame_button.isEnabled() + assert controls.reset_all_button.isEnabled() + viewer.dims.order = (1, 0, 2) + qtbot.waitUntil(lambda: not controls.reset_frame_button.isEnabled()) + assert controls.reset_all_button.isEnabled() + viewer.dims.order = (0, 1, 2) + qtbot.waitUntil(controls.reset_frame_button.isEnabled) + points.data[1, 1:] = [100, 200] + points.events.data(action="changed", data_indices=(1,)) + points.remove([2]) + viewer.dims.set_current_step(0, 1) + controls.reset_frame_button.click() + np.testing.assert_array_equal(tracks.data[1], original[1]) + assert len(points.data) == 3 + controls.reset_all_button.click() + np.testing.assert_array_equal(tracks.data, original) + np.testing.assert_array_equal(timeline._edited_frames, [0]) + assert timeline._removed_points == [] + viewer.layers.remove(tracks) + qtbot.waitUntil(lambda: not controls.reset_all_button.isEnabled()) + + +def test_reset_controls_without_loaded_pose(qtbot): + """Reset is unavailable for unrelated points or an empty viewer.""" + from movement.napari.edit_timeline_widget import EditControlsWidget + + viewer = ViewerModel() + controls = EditControlsWidget(napari_viewer=viewer) + qtbot.addWidget(controls) + assert not controls.reset_all_button.isEnabled() + viewer.add_points([[0, 1, 2]]) + controls._update_reset_enabled() + assert not controls.reset_frame_button.isEnabled() + assert not controls.reset_all_button.isEnabled() + + +@pytest.mark.parametrize("answer", [QMessageBox.Reset, QMessageBox.Cancel]) +def test_reset_all_requires_confirmation( + reset_layers, qtbot, monkeypatch, answer +): + """Only explicit confirmation restores points; cancel keeps edits.""" + from unittest.mock import Mock + + from movement.napari.edit_timeline_widget import EditControlsWidget + + points, tracks, original, props, viewer = reset_layers + controls = EditControlsWidget(napari_viewer=viewer) + qtbot.addWidget(controls) + points.remove([0]) + edited = tracks.data.copy() + warning = Mock(return_value=answer) + monkeypatch.setattr(QMessageBox, "warning", warning) + controls.reset_all_button.click() + warning.assert_called_once() + assert warning.call_args.kwargs["defaultButton"] == QMessageBox.Cancel + expected = original if answer == QMessageBox.Reset else edited + np.testing.assert_array_equal(tracks.data, expected) + + +def test_loaded_state_roundtrip(valid_poses_dataset, tmp_path, qtbot): + """Loading stores a state that restores missing data on export.""" + import xarray as xr + + from movement.napari.convert import napari_layers_to_ds + from movement.napari.loader_widgets import DataLoader + + ds = valid_poses_dataset.copy(deep=True) + ds.position.values.flat[0:2] = np.nan + path = tmp_path / "loaded.nc" + ds.to_netcdf(path) + viewer = ViewerModel() + loader = DataLoader(viewer) + qtbot.addWidget(loader) + loader.source_software_combo.setCurrentText("movement (netCDF)") + loader.file_path_edit.setText(str(path)) + loader._on_load_clicked() + points = loader.points_layer + assert wiring.POINTS_POSITION_KEY in points.metadata + loaded = napari_layers_to_ds( + points.data, + points.properties, + points.metadata[wiring.POINTS_PROPERTIES_KEY], + ds.attrs, + ) + points.data[0, 1:] += 50 + points.events.data(action="changed", data_indices=(0,)) + points.remove([1]) + wiring.reset_edits(points) + restored = napari_layers_to_ds( + points.data, + points.properties, + points.metadata[wiring.POINTS_PROPERTIES_KEY], + ds.attrs, + ) + xr.testing.assert_allclose(restored.position, loaded.position) + xr.testing.assert_allclose(restored.confidence, loaded.confidence) + + +def test_reset_in_docked_widget( + make_napari_viewer_proxy, + valid_poses_path_and_ds, + loaded_data_loader, + qtbot, + monkeypatch, +): + """Exercise reset buttons through the real viewer and docked timeline.""" + from movement.napari.loader_widgets import DataLoader + from movement.napari.meta_widget import MovementMetaWidget + + monkeypatch.setattr( + QMessageBox, "warning", lambda *a, **kw: QMessageBox.Reset + ) + viewer = make_napari_viewer_proxy() + widget = MovementMetaWidget(viewer) + qtbot.addWidget(widget) + loader = widget.findChild(DataLoader) + path, ds = valid_poses_path_and_ds + loaded_data_loader(path, ds, loader=loader) + points = loader.points_layer + tracks = points.metadata[wiring.TRACKS_LAYER_KEY] + original = tracks.data.copy() + points.data[0, 1:] += 50 + points.events.data(action="changed", data_indices=(0,)) + qtbot.waitUntil(widget.edit_controls.reset_frame_button.isEnabled) + assert widget.edit_timeline_widget is not None + viewer.dims.set_current_step(0, int(original[0, 1])) + widget.edit_controls.reset_frame_button.click() + np.testing.assert_array_equal(tracks.data, original) + points.remove([0]) + widget.edit_controls.reset_all_button.click() + np.testing.assert_array_equal(tracks.data, original) + assert widget.edit_timeline_widget._removed_points == []