70 replace tracks controller

src/funtracks/actions/__init__.py

@@ -0,0 +1,6 @@
+from ._base import ActionGroup, TracksAction
+from .add_delete_edge import AddEdge, DeleteEdge
+from .add_delete_node import AddNode, DeleteNode
+from .update_node_attrs import UpdateNodeAttrs
+from .update_segmentation import UpdateNodeSeg
+from .update_track_id import UpdateTrackID

src/funtracks/actions/_base.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from typing_extensions import override
+
+if TYPE_CHECKING:
+    from funtracks.data_model import SolutionTracks
+
+
+class TracksAction:
+    def __init__(self, tracks: SolutionTracks):
+        """An modular change that can be applied to the given Tracks. The tracks must
+        be passed in at construction time so that metadata needed to invert the action
+        can be extracted.
+        The change should be applied in the init function.
+
+        Args:
+            tracks (Tracks): The tracks that this action will edit
+        """
+        self.tracks = tracks
+
+    def inverse(self) -> TracksAction:
+        """Get the inverse of this action. Calling this function does undo the action,
+        since the change is applied in the action constructor.
+
+        Raises:
+            NotImplementedError: if the inverse is not implemented in the subclass
+
+        Returns:
+            TracksAction: An action that un-does this action, bringing the tracks
+                back to the exact state it had before applying this action.
+        """
+        raise NotImplementedError("Inverse not implemented")
+
+
+class ActionGroup(TracksAction):
+    def __init__(
+        self,
+        tracks: SolutionTracks,
+        actions: list[TracksAction],
+    ):
+        """A group of actions that is also an action, used to modify the given tracks.
+        This is useful for creating composite actions from the low-level actions.
+        Composite actions can contain application logic and can be un-done as a group.
+
+        Args:
+            tracks (Tracks): The tracks that this action will edit
+            actions (list[TracksAction]): A list of actions contained within the group,
+                in the order in which they should be executed.
+        """
+        super().__init__(tracks)
+        self.actions = actions
+
+    @override
+    def inverse(self) -> ActionGroup:
+        actions = [action.inverse() for action in self.actions[::-1]]
+        return ActionGroup(self.tracks, actions)

src/funtracks/data_model/action_history.py → src/funtracks/actions/action_history.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from .actions import TracksAction  # noqa
+    from ._base import TracksAction  # noqa
 
 
 class ActionHistory:

src/funtracks/actions/add_delete_edge.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from ._base import TracksAction
+
+if TYPE_CHECKING:
+    from funtracks.data_model import SolutionTracks
+    from funtracks.data_model.tracks import Edge
+
+
+class AddEdge(TracksAction):
+    """Action for adding new edges"""
+
+    def __init__(self, tracks: SolutionTracks, edge: Edge):
+        super().__init__(tracks)
+        self.edge = edge
+        self._apply()
+
+    def inverse(self) -> TracksAction:
+        """Delete edges"""
+        return DeleteEdge(self.tracks, self.edge)
+
+    def _apply(self) -> None:
+        """
+        Steps:
+        - add each edge to the graph. Assumes all edges are valid (they should be checked
+        at this point already)
+        """
+        attrs: dict[str, Sequence[Any]] = {}
+        attrs.update(self.tracks._compute_edge_attrs(self.edge))
+        for node in self.edge:
+            if not self.tracks.graph.has_node(node):
+                raise ValueError(
+                    f"Cannot add edge {self.edge}: endpoint {node} not in graph yet"
+                )
+        self.tracks.graph.add_edge(self.edge[0], self.edge[1], **attrs)
+
+
+class DeleteEdge(TracksAction):
+    """Action for deleting edges"""
+
+    def __init__(self, tracks: SolutionTracks, edge: Edge):
+        super().__init__(tracks)
+        self.edge = edge
+        self._apply()
+
+    def inverse(self) -> TracksAction:
+        """Restore edges and their attributes"""
+        return AddEdge(self.tracks, self.edge)
+
+    def _apply(self) -> None:
+        """Steps:
+        - Remove the edges from the graph
+        """
+        if self.tracks.graph.has_edge(*self.edge):
+            self.tracks.graph.remove_edge(*self.edge)
+        else:
+            raise ValueError(f"Edge {self.edge} not in the graph, and cannot be removed")

src/funtracks/actions/add_delete_node.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from funtracks.data_model.graph_attributes import NodeAttr
+from funtracks.data_model.solution_tracks import SolutionTracks
+
+from ._base import TracksAction
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from funtracks.data_model import SolutionTracks
+    from funtracks.data_model.tracks import Node, SegMask
+
+
+class AddNode(TracksAction):
+    """Action for adding new nodes. If a segmentation should also be added, the
+    pixels for each node should be provided. The label to set the pixels will
+    be taken from the node id. The existing pixel values are assumed to be
+    zero - you must explicitly update any other segmentations that were overwritten
+    using an UpdateNodes action if you want to be able to undo the action.
+    """
+
+    def __init__(
+        self,
+        tracks: SolutionTracks,
+        node: Node,
+        attributes: dict[str, Any],
+        pixels: SegMask | None = None,
+    ):
+        """Create an action to add a new node, with optional segmentation
+
+        Args:
+            tracks (Tracks): The Tracks to add the node to
+            node (Node): A node id
+            attributes (Attrs): Includes times, track_ids, and optionally positions
+            pixels (SegMask | None, optional): The segmentation associated with
+                the node. Defaults to None.
+        """
+        super().__init__(tracks)
+        self.node = node
+        user_attrs = attributes.copy()
+        if NodeAttr.TIME.value not in attributes:
+            raise ValueError("Must provide a time attribute for each node")
+        self.time = attributes.pop(NodeAttr.TIME.value)
+        self.position = attributes.pop(NodeAttr.POS.value, None)
+        self.pixels = pixels
+        self.attributes = user_attrs
+        self._apply()
+
+    def inverse(self) -> TracksAction:
+        """Invert the action to delete nodes instead"""
+        return DeleteNode(self.tracks, self.node)
+
+    def _apply(self) -> None:
+        """Apply the action, and set segmentation if provided in self.pixels"""
+        if self.pixels is not None:
+            self.tracks.set_pixels(self.pixels, self.node)
+        attrs = self.attributes
+        self.tracks.graph.add_node(self.node)
+        self.tracks.set_time(self.node, self.time)
+        final_pos: np.ndarray
+        if self.tracks.segmentation is not None:
+            computed_attrs = self.tracks._compute_node_attrs(self.node, self.time)
+            if self.position is None:
+                final_pos = np.array(computed_attrs[NodeAttr.POS.value])
+            else:
+                final_pos = self.position
+            attrs[NodeAttr.AREA.value] = computed_attrs[NodeAttr.AREA.value]
+        elif self.position is None:
+            raise ValueError("Must provide positions or segmentation and ids")
+        else:
+            final_pos = self.position
+
+        self.tracks.set_position(self.node, final_pos)
+        for attr, values in attrs.items():
+            self.tracks._set_node_attr(self.node, attr, values)
+
+        track_id = attrs[NodeAttr.TRACK_ID.value]
+        if track_id not in self.tracks.track_id_to_node:
+            self.tracks.track_id_to_node[track_id] = []
+        self.tracks.track_id_to_node[track_id].append(self.node)
+
+
+class DeleteNode(TracksAction):
+    """Action of deleting existing nodes
+    If the tracks contain a segmentation, this action also constructs a reversible
+    operation for setting involved pixels to zero
+    """
+
+    def __init__(
+        self,
+        tracks: SolutionTracks,
+        node: Node,
+        pixels: SegMask | None = None,
+    ):
+        super().__init__(tracks)
+        self.node = node
+        self.attributes = {
+            NodeAttr.TIME.value: self.tracks.get_time(node),
+            NodeAttr.POS.value: self.tracks.get_position(node),
+            NodeAttr.TRACK_ID.value: self.tracks.get_node_attr(
+                node, NodeAttr.TRACK_ID.value
+            ),
+        }
+        self.pixels = self.tracks.get_pixels(node) if pixels is None else pixels
+        self._apply()
+
+    def inverse(self) -> TracksAction:
+        """Invert this action, and provide inverse segmentation operation if given"""
+
+        return AddNode(self.tracks, self.node, self.attributes, pixels=self.pixels)
+
+    def _apply(self) -> None:
+        """ASSUMES THERE ARE NO INCIDENT EDGES - raises valueerror if an edge will be
+        removed by this operation
+        Steps:
+        - For each node
+            set pixels to 0 if self.pixels is provided
+        - Remove nodes from graph
+        """
+        if self.pixels is not None:
+            self.tracks.set_pixels(self.pixels, 0)
+
+        if isinstance(self.tracks, SolutionTracks):
+            self.tracks.track_id_to_node[self.tracks.get_track_id(self.node)].remove(
+                self.node
+            )
+
+        self.tracks.graph.remove_node(self.node)

src/funtracks/actions/update_node_attrs.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from funtracks.data_model.graph_attributes import NodeAttr
+
+from ._base import TracksAction
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from funtracks.data_model import SolutionTracks
+    from funtracks.data_model.tracks import Node
+
+
+class UpdateNodeAttrs(TracksAction):
+    """Action for user updates to node attributes. Cannot update protected
+    attributes (time, area, track id), as these are controlled by internal application
+    logic."""
+
+    def __init__(
+        self,
+        tracks: SolutionTracks,
+        node: Node,
+        attrs: dict[str, Any],
+    ):
+        """
+        Args:
+            tracks (Tracks): The tracks to update the node attributes for
+            node (Node): The node to update the attributes for
+            attrs (dict[str, Any]): A mapping from attribute name to list of new attribute
+                values for the given nodes.
+
+        Raises:
+            ValueError: If a protected attribute is in the given attribute mapping.
+        """
+        super().__init__(tracks)
+        protected_attrs = [
+            tracks.time_attr,
+            NodeAttr.AREA.value,
+            NodeAttr.TRACK_ID.value,
+        ]
+        for attr in attrs:
+            if attr in protected_attrs:
+                raise ValueError(f"Cannot update attribute {attr} manually")
+        self.node = node
+        self.prev_attrs = {attr: self.tracks.get_node_attr(node, attr) for attr in attrs}
+        self.new_attrs = attrs
+        self._apply()
+
+    def inverse(self) -> TracksAction:
+        """Restore previous attributes"""
+        return UpdateNodeAttrs(
+            self.tracks,
+            self.node,
+            self.prev_attrs,
+        )
+
+    def _apply(self) -> None:
+        """Set new attributes"""
+        for attr, value in self.new_attrs.items():
+            self.tracks._set_node_attr(self.node, attr, value)