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,17 @@
 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, QPushButton, QVBoxLayout, QWidget
 
 from movement.napari.layer_wiring import (
     MAX_FRAME_IDX_KEY,
+    POINTS_BASELINE_KEY,
     POINTS_PROPERTIES_KEY,
+    TRACKS_LAYER_KEY,
     active_movement_points_layer,
+    frame_axis_is_sliced,
     is_movement_points_layer,
+    reset_points_to_baseline,
 )
 
 if TYPE_CHECKING:
@@ -51,9 +55,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 +78,65 @@ 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_BASELINE_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 and their baseline 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_points_to_baseline(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:
+            reset_points_to_baseline(layer)
+
 
 class EditTimelineWidget(QWidget):
     """Dock widget flagging frames with edited points.
@@ -283,6 +345,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

@@ -38,6 +38,7 @@
 DATASET_ATTRS_KEY: str = "movement_dataset_attrs"
 TRACKS_LAYER_KEY: str = "movement_tracks_layer"
 MAX_FRAME_IDX_KEY: str = "movement_max_frame_idx"
+POINTS_BASELINE_KEY: str = "movement_points_baseline"
 
 # Keep a set of viewers already wired by connect_viewer_callbacks,
 # so we don't wire them twice. We use a WeakSet so tracking a viewer here
@@ -186,6 +187,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 +270,89 @@ def set_tracks_layer_data(tracks_layer, data, properties):
         tracks_layer.data = data
         tracks_layer.properties = properties
         tracks_layer.color_by = color_by
+
+
+def capture_points_baseline(points_layer: Points) -> None:
+    """Store independent copies of pose data immediately after loading.
+
+    The baseline includes edits saved in the input file. It is owned by the
+    layer, not the loader widget, and is never updated by subsequent edits.
+    """
+    tracks = points_layer.metadata[TRACKS_LAYER_KEY]
+    points_layer.metadata[POINTS_BASELINE_KEY] = {
+        "data": tracks.data.copy(),
+        "points_properties": {
+            key: value.copy() for key, value in points_layer.properties.items()
+        },
+        "tracks_properties": {
+            key: value.copy() for key, value in tracks.properties.items()
+        },
+        "symbol": points_layer.symbol.copy(),
+    }
+
+
+def reset_points_to_baseline(
+    points_layer: Points, frame: int | None = None
+) -> None:
+    """Restore loaded pose points for one frame, or all frames if None.
+
+    Deleted points are restored from the baseline, rather than inferred
+    from surviving edited flags. Other frames retain their current data.
+    Identity swaps and individual edit history are not supported.
+    """
+    baseline = points_layer.metadata[POINTS_BASELINE_KEY]
+    tracks = points_layer.metadata[TRACKS_LAYER_KEY]
+    restore = (
+        baseline["data"][:, 1] == frame
+        if frame is not None
+        else np.ones(len(baseline["data"]), dtype=bool)
+    )
+    keep = (
+        tracks.data[:, 1] != frame
+        if frame is not None
+        else np.zeros(len(tracks.data), dtype=bool)
+    )
+    data = np.concatenate([tracks.data[keep], baseline["data"][restore]])
+    order = np.lexsort((data[:, 1], data[:, 0]))
+    data = data[order]
+
+    def merge_properties(current, original):
+        keys = current.keys() | original.keys()
+        return {
+            key: np.concatenate(
+                [
+                    current.get(key, np.zeros(len(tracks.data), dtype=bool))[
+                        keep
+                    ],
+                    original.get(
+                        key, np.zeros(len(baseline["data"]), dtype=bool)
+                    )[restore],
+                ]
+            )[order]
+            for key in keys
+        }
+
+    properties = merge_properties(
+        points_layer.properties, baseline["points_properties"]
+    )
+    track_properties = merge_properties(
+        tracks.properties, baseline["tracks_properties"]
+    )
+    symbols = np.concatenate(
+        [points_layer.symbol[keep], baseline["symbol"][restore]]
+    )[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

@@ -30,6 +30,7 @@
     POINTS_LAYER_KEY,
     POINTS_PROPERTIES_KEY,
     TRACKS_LAYER_KEY,
+    capture_points_baseline,
     connect_viewer_callbacks,
     frame_axis_is_sliced,
     on_points_data_changed,
@@ -212,6 +213,8 @@ def _on_load_clicked(self):
         # and a boxes layer if the dataset is a bounding boxes one
         self._add_points_layer()
         self._add_tracks_layer()
+        if self.data_bboxes is None:
+            capture_points_baseline(self.points_layer)
         if self.data_bboxes is not None:
             self._add_boxes_layer()
 

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
         )