Add frame and full pose reset controls

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.
 

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()

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,
+    )

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)

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
         )