"""Layer callbacks that must outlive the widget connecting them.
Put a callback here if it should be active as long as the viewer
or a layer exists (rather than sharing lifetime with a widget that
can be closed).
- Typical candidates are callbacks that mutate layer state (data, properties,
symbols, ``editable``): that state is part of the user's work and may be
saved to a file later, so it must survive the widget being closed.
- A callback that only refreshes a widget's own UI (a dropdown, a table, a
button) can stay a method on that widget: once the widget is gone there is
nothing left to update, which is fine.
"""
import warnings
from functools import partial
from weakref import WeakKeyDictionary, 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 napari.utils.events import Event
from movement.napari.convert import LoadedPointKeys, removed_points_mask
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.
#
# Poses Points layers also carry an extra ``events.edits_reset`` event
# (see add_edits_reset_event), emitted by reset_edits once Points and
# Tracks are consistent again. Its ``frame`` attribute is the reset frame,
# or None if all frames were reset.
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"
# ---- Layer helpers -------------------------------------
[docs]
def is_movement_points_layer(layer) -> bool:
"""Return ``True`` for a movement-loaded napari Points layer."""
layer = getattr(layer, "__wrapped__", layer)
return isinstance(layer, Points) and bool(
layer.metadata.get(POINTS_LAYER_KEY)
)
[docs]
def active_movement_points_layer(viewer):
"""Return the active movement Points layer, else the last one, else None.
Prefers napari's active selection, but falls back to scanning the layer
list so a Points layer is still found when the active selection is
something else or unset.
"""
active = viewer.layers.selection.active
if is_movement_points_layer(active):
return getattr(active, "__wrapped__", active)
for layer in reversed(viewer.layers):
if is_movement_points_layer(layer):
return getattr(layer, "__wrapped__", layer)
return None
[docs]
def drop_internal_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"),
]
# Per-layer cache of the loaded point keys used by removed_points
_LOADED_POINT_KEYS: WeakKeyDictionary = WeakKeyDictionary()
[docs]
def removed_points(layer: Points) -> tuple[np.ndarray, np.ndarray]:
"""Return the frames and individuals of points removed from ``layer``.
Removal is defined by :func:`~movement.napari.convert.removed_points_mask`,
the same rule used to flag points as ``edited`` on save.
Parameters
----------
layer
A movement Points layer. Its loaded state is read from
``layer.metadata[POINTS_PROPERTIES_KEY]``.
Returns
-------
frames : numpy.ndarray
Frame index of each removed point.
individuals : numpy.ndarray
Individual of each removed point, aligned with ``frames``.
Both arrays are empty if the layer was not created by ``movement``.
"""
loaded = layer.metadata.get(POINTS_PROPERTIES_KEY)
if loaded is None:
return np.array([], dtype=int), np.array([], dtype=object)
cached = _LOADED_POINT_KEYS.get(layer)
if cached is not None and cached[0] is loaded:
keys = cached[1]
else:
keys = LoadedPointKeys.from_properties(loaded)
_LOADED_POINT_KEYS[layer] = (loaded, keys)
removed = removed_points_mask(
loaded, layer.data[:, 0], layer.features, keys
)
return keys.frames[removed], loaded["individual"].to_numpy()[removed]
# ---- Callbacks with viewer lifetime --------------------
# 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
# does not prevent it from being garbage-collected when it is no longer used.
_WIRED_VIEWERS: WeakSet = WeakSet()
[docs]
def connect_viewer_callbacks(viewer) -> None:
"""Wire the layer callbacks to a viewer, skipping if already wired.
These wirings last as long as the viewer, no matter which widget requested
them.
"""
# Check if viewer has been wired by this function
viewer = getattr(viewer, "__wrapped__", viewer)
if viewer in _WIRED_VIEWERS:
return
# Connect relevant layer callbacks to the viewer.
# Connect frame slider range update to layer events
for action in (
"inserted",
"removed",
):
getattr(viewer.layers.events, action).connect(
partial(update_frame_slider_range, viewer)
)
# Point drags are only guaranteed to stay within their own frame
# when frame is the sliced (non-displayed) axis in a 2D view. If
# axes are rolled or a 3D view is used in the viewer, disable editing
# rather than risk a drag moving a point onto a different frame.
for event in (
viewer.dims.events.order,
viewer.dims.events.ndisplay,
):
event.connect(partial(update_points_layers_editable, viewer))
# Update set
_WIRED_VIEWERS.add(viewer)
[docs]
def update_frame_slider_range(viewer, event=None):
"""Widen the frame slider range to cover NaN-trimmed frames.
napari derives ``viewer.dims.range`` from the world-coordinate union of
all layer extents, and it does so before this callback runs. Movement's
layers have their leading/trailing all-NaN rows dropped, so a dataset
that starts or ends with NaNs yields an extent — and therefore a slider
range — narrower than the real frame span.
Extend napari's range to cover the true span of every movement layer,
but never replace it: layers movement did not create (a video the user
opened, another plugin's layer) keep the range napari computed for them,
including any scale or translate they carry.
"""
max_frame_indices = [
ly.metadata[MAX_FRAME_IDX_KEY]
for ly in viewer.layers
if MAX_FRAME_IDX_KEY in getattr(ly, "metadata", {})
]
if not max_frame_indices:
return
current_range = viewer.dims.range[0]
start = min(current_range.start, 0.0)
stop = max(current_range.stop, *max_frame_indices)
if (start, stop) != (current_range.start, current_range.stop):
viewer.dims.range = (
RangeTuple(start=start, stop=stop, step=current_range.step),
) + viewer.dims.range[1:]
[docs]
def frame_axis_is_sliced(viewer) -> bool:
"""Determine whether the frame axis is the sliced axis in a 2D view."""
return viewer.dims.ndisplay == 2 and viewer.dims.order[0] == 0
[docs]
def update_points_layers_editable(viewer, event=None):
"""Disable point editing while the frame axis isn't sliced.
Connected to ``viewer.dims.events.order``/``ndisplay``.
It disables editing on every movement Points layer if
the viewer axes are rolled or switched to 3D; napari greys
out the select/add/delete controls.
"""
is_editable = frame_axis_is_sliced(viewer)
for layer in viewer.layers:
if is_movement_points_layer(layer):
layer.editable = is_editable
# ---- Callbacks with layer lifetime --------------------
[docs]
def set_point_symbol_by_edited(layer: Points) -> None:
"""Show points flagged as edited with a distinct marker symbol."""
edited = layer.properties.get("edited")
if edited is None or not edited.any():
return
symbols = np.asarray(layer.symbol).copy()
symbols[edited] = EDITED_POINT_SYMBOL
layer.symbol = symbols
[docs]
def on_points_data_changed(event):
"""Keep the corresponding Tracks layer in sync with the Points layer.
Connected to ``points_layer.events.data``. Handles two actions:
- ``ActionType.CHANGED`` (a point was dragged): sets the
confidence score of moved points to NaN, marks them as
edited, and changes their marker symbol to
``EDITED_POINT_SYMBOL`` so edited points are visually
distinguishable. The Tracks layer row is updated in place
via `sync_tracks_layer`.
- ``ActionType.REMOVED`` (one or more points were deleted):
removes the corresponding rows from
the Tracks layer via `remove_from_tracks_layer`.
"""
layer = event.source
if not isinstance(layer, Points):
return
if event.action == ActionType.CHANGED:
moved_indices = list(event.data_indices)
props = layer.properties
props["confidence"] = props["confidence"].copy()
props["confidence"][moved_indices] = float("nan")
if "edited" in props:
props["edited"] = props["edited"].copy()
else:
props["edited"] = np.full(len(props["confidence"]), False)
props["edited"][moved_indices] = True
layer.properties = props
set_point_symbol_by_edited(layer)
sync_tracks_layer(layer, moved_indices)
elif event.action == ActionType.REMOVED:
removed_indices = list(event.data_indices)
remove_from_tracks_layer(layer, removed_indices)
[docs]
def sync_tracks_layer(points_layer, moved_indices):
"""Update the corresponding Tracks layer to match an edited point.
A moved point's new (frame, y, x) is written to the same row
in the Tracks layer, so the track segment connecting the
previous frame to this one terminates at the dragged position.
"""
tracks_layer = points_layer.metadata[TRACKS_LAYER_KEY]
# Points and Tracks layers are built from the same NaN-filtered
# array in the same row order (see _add_points_layer/
# _add_tracks_layer). The Tracks layer only has an extra
# leading track_id column.
tracks_data = tracks_layer.data
tracks_data[moved_indices, 1:] = points_layer.data[moved_indices]
set_tracks_layer_data(tracks_layer, tracks_data, tracks_layer.properties)
[docs]
def remove_from_tracks_layer(points_layer, removed_indices):
"""Remove the rows corresponding to deleted points.
Users edit the Points layer directly, either dragging points
or removing inaccurate predictions. The Tracks layer has no
interactive editing of its own, so it must be kept in sync
with the Points layer instead.
``removed_indices`` are indices in the Points layer which line
up with rows in the Tracks layer, the same way
:func:`sync_tracks_layer` relies on for edits.
"""
tracks_layer = points_layer.metadata[TRACKS_LAYER_KEY]
tracks_data = np.delete(tracks_layer.data, removed_indices, axis=0)
tracks_properties = {
key: np.delete(np.asarray(values), removed_indices, axis=0)
for key, values in tracks_layer.properties.items()
}
set_tracks_layer_data(tracks_layer, tracks_data, tracks_properties)
[docs]
def set_tracks_layer_data(tracks_layer, data, properties):
"""Set a Tracks layer's data and properties, preserving color_by.
Setting ``.data`` on a napari Tracks layer resets its internal
features to empty, which transiently invalidates ``color_by``
(napari warns and falls back to "track_id") even though we
restore the same properties and colour-by property right
after. Suppress that spurious warning around the sequence.
"""
color_by = tracks_layer.color_by
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message=".*Previous color_by key.*",
category=UserWarning,
)
tracks_layer.data = data
tracks_layer.properties = properties
tracks_layer.color_by = color_by
[docs]
def add_edits_reset_event(points_layer: Points) -> None:
"""Register the ``edits_reset`` event emitted by `reset_edits`.
Must be called before the layer is added to the viewer, so that
widgets tracking newly inserted layers can connect to it.
"""
points_layer.events.add(edits_reset=Event)
[docs]
def reset_edits(points_layer: Points, frame: int | None = None) -> None:
"""Restore loaded 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.
"""
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[key])[keep]
if key in current
else np.zeros(keep.sum(), bool),
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, drop_internal_properties(loaded_props)
)
track_properties = merge_properties(tracks.properties, loaded_props)
# napari's auto-generated track_id is absent from the loaded props,
# so rebuild it from the merged data instead of zero-filling it
track_properties["track_id"] = data[:, 0].astype(int)
# 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, via a dedicated event, 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.edits_reset(frame=frame)
