"""Load pose tracking data from various frameworks into ``movement``."""
import warnings
from pathlib import Path
from typing import Literal, cast
import h5py
import numpy as np
import pandas as pd
import pynwb
import xarray as xr
from sleap_io.io.slp import read_labels
from sleap_io.model.labels import Labels
from movement.io.load import register_loader
from movement.utils.logging import logger
from movement.validators.datasets import ValidPosesInputs
from movement.validators.files import (
ValidAniposeCSV,
ValidCOCOKeypointResults,
ValidDeepLabCutCSV,
ValidDeepLabCutH5,
ValidFile,
ValidNWBFile,
ValidSleapAnalysis,
ValidSleapLabels,
)
[docs]
def from_numpy(
position_array: np.ndarray,
confidence_array: np.ndarray | None = None,
individual_names: list[str] | None = None,
keypoint_names: list[str] | None = None,
frame_array: np.ndarray | None = None,
fps: float | None = None,
source_software: str | None = None,
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from NumPy arrays.
Parameters
----------
position_array
Array of shape (n_frames, n_space, n_keypoints, n_individuals)
containing the poses. It will be converted to a
:class:`xarray.DataArray` object named "position".
confidence_array
Array of shape (n_frames, n_keypoints, n_individuals) containing
the point-wise confidence scores. It will be converted to a
:class:`xarray.DataArray` object named "confidence".
If None (default), the scores will be set to an array of NaNs.
individual_names
List of unique names for the individuals in the video. If None
(default), the individuals will be named "id_0", "id_1", etc.
keypoint_names
List of unique names for the keypoints in the skeleton. If None
(default), the keypoints will be named "keypoint_0", "keypoint_1",
etc.
frame_array
Array containing the frame numbers corresponding to the data.
It must be a column vector of shape (n_frames, 1) and values
must be monotonically increasing. If None (default), frame
numbers are assigned as consecutive 0-based integers.
fps
Frames per second of the video. Defaults to None, in which case
the time coordinates will be in frame numbers.
source_software
Name of the pose estimation software from which the data originate.
Defaults to None.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
Examples
--------
Create random position data for two individuals with three keypoints
each, tracked in 2D space for 100 frames. As confidence scores are omitted,
they default to NaN, and individuals/keypoints are auto-named
("id_0", "id_1", ... and "keypoint_0", "keypoint_1", ...). Time
coordinates default to consecutive 0-based frame numbers.
>>> import numpy as np
>>> from movement.io import load_poses
>>> rng = np.random.default_rng(seed=42)
>>> ds = load_poses.from_numpy(
... position_array=rng.random((100, 2, 3, 2)),
... )
Create a dataset with the same data as above, but name the individuals
``Alice`` and ``Bob``, the keypoints ``snout``, ``centre``, and
``tail_base``, and set all confidence scores to 1.
>>> ds = load_poses.from_numpy(
... position_array=rng.random((100, 2, 3, 2)),
... confidence_array=np.ones((100, 3, 2)),
... individual_names=["Alice", "Bob"],
... keypoint_names=["snout", "centre", "tail_base"],
... )
Create a dataset with the same data as above, but specify that the
100 frames are numbered from the start frame 1200 to the end frame
1299, instead of defaulting to 0-based frame numbers.
>>> ds = load_poses.from_numpy(
... position_array=rng.random((100, 2, 3, 2)),
... confidence_array=np.ones((100, 3, 2)),
... individual_names=["Alice", "Bob"],
... keypoint_names=["snout", "centre", "tail_base"],
... frame_array=np.arange(1200, 1300).reshape(-1, 1),
... )
Create a dataset with the same data as above, but with the time
coordinates in seconds. We use a video sampling rate of 30 fps. The time
coordinates in the resulting dataset will indicate the elapsed time from
the capture of the 0th frame. So for the frames 1200, 1201, 1202, ...
1299 the corresponding time coordinates in seconds will be
40.0, 40.033..., 40.067..., ... 43.3 s.
>>> ds = load_poses.from_numpy(
... position_array=rng.random((100, 2, 3, 2)),
... confidence_array=np.ones((100, 3, 2)),
... individual_names=["Alice", "Bob"],
... keypoint_names=["snout", "centre", "tail_base"],
... frame_array=np.arange(1200, 1300).reshape(-1, 1),
... fps=30,
... )
"""
valid_poses_inputs = ValidPosesInputs(
position_array=position_array,
confidence_array=confidence_array,
individual_names=individual_names,
keypoint_names=keypoint_names,
frame_array=frame_array,
fps=fps,
source_software=source_software,
)
return valid_poses_inputs.to_dataset()
[docs]
def from_dlc_style_df(
df: pd.DataFrame,
fps: float | None = None,
source_software: Literal["DeepLabCut", "LightningPose"] = "DeepLabCut",
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a DeepLabCut-style DataFrame.
Parameters
----------
df
DataFrame containing the pose tracks and confidence scores. Must
be formatted as in DeepLabCut output files (see Notes).
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame numbers.
source_software
Name of the pose estimation software from which the data originate.
Defaults to "DeepLabCut", but it can also be "LightningPose"
(because they use the same DataFrame format).
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
Notes
-----
The DataFrame must have a multi-index column with the following levels:
"scorer", ("individuals"), "bodyparts", "coords".
The "individuals" level may be omitted if there is only one individual
in the video.
The "coords" level contains either:
- the spatial coordinates "x", "y", and "likelihood"
(point-wise confidence scores), or
- the spatial coordinates "x", "y", and "z" (3D poses estimated by
`triangulating 2D poses from multiple DeepLabCut output files\
<https://deeplabcut.github.io/DeepLabCut/docs/Overviewof3D.html>`__).
The row index corresponds to the frame number.
See Also
--------
movement.io.load_poses.from_dlc_file
"""
# Read names of individuals and keypoints from the DataFrame
individual_names = (
df.columns.get_level_values("individuals").unique().to_list()
if "individuals" in df.columns.names
else None
)
keypoint_names = (
df.columns.get_level_values("bodyparts").unique().to_list()
)
# Extract position (and confidence if present)
coord_names = df.columns.get_level_values("coords").unique().to_list()
n_coords = len(coord_names)
n_individuals = len(individual_names) if individual_names else 1
tracks = (
df.to_numpy()
.reshape((-1, n_individuals, len(keypoint_names), n_coords))
.transpose(0, 3, 2, 1)
)
if "likelihood" in coord_names: # Coords: ['x', 'y', 'likelihood']
likelihood_index = coord_names.index("likelihood")
confidence_array = tracks[:, likelihood_index, :, :]
pos_idx = [j for j in range(n_coords) if j != likelihood_index]
position_array = tracks[:, pos_idx, :, :]
else: # Coords: ['x', 'y', 'z']
position_array = tracks
confidence_array = None
return from_numpy(
position_array=position_array,
confidence_array=confidence_array,
individual_names=individual_names,
keypoint_names=keypoint_names,
fps=fps,
source_software=source_software,
)
[docs]
@register_loader(
"SLEAP", file_validators=[ValidSleapLabels, ValidSleapAnalysis]
)
def from_sleap_file(file: str | Path, fps: float | None = None) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a SLEAP file.
Parameters
----------
file
Path to the file containing the SLEAP predictions in .h5
(analysis) format. Alternatively, a .slp (labels) file can
also be supplied (but this feature is experimental, see Notes).
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame numbers.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
Notes
-----
The SLEAP predictions are normally saved in .slp files, e.g.
"v1.predictions.slp". An analysis file, suffixed with ".h5" can be exported
from the .slp file, using either the command line tool `sleap-convert`
(with the "--format analysis" option enabled) or the SLEAP GUI (Choose
"Export Analysis HDF5…" from the "File" menu) [1]_. This is the
preferred format for loading pose tracks from SLEAP into *movement*.
You can also directly load the .slp file. However, if the file contains
multiple videos, only the pose tracks from the first video will be loaded.
If the file contains a mix of user-labelled and predicted instances, user
labels are prioritised over predicted instances to mirror SLEAP's approach
when exporting .h5 analysis files [2]_.
*movement* expects the tracks to be assigned and proofread before loading
them, meaning each track is interpreted as a single individual. If
no tracks are found in the file, *movement* assumes that this is a
single-individual track, and will assign a default individual name.
If multiple instances without tracks are present in a frame, the last
instance is selected [2]_.
Follow the SLEAP guide for tracking and proofreading [3]_.
References
----------
.. [1] https://docs.sleap.ai/latest/tutorial/exporting-the-results/#analysis-hdf5
.. [2] https://github.com/talmolab/sleap/blob/v1.3.3/sleap/info/write_tracking_h5.py#L59
.. [3] https://docs.sleap.ai/latest/guides/tracking-and-proofreading/
Examples
--------
>>> from movement.io import load_poses
>>> ds = load_poses.from_sleap_file("path/to/file.analysis.h5", fps=30)
"""
valid_file = cast("ValidFile", file)
file_path = valid_file.file
if file_path.suffix == ".h5":
ds = _ds_from_sleap_analysis_file(file_path, fps=fps)
else: # file.suffix == ".slp"
ds = _ds_from_sleap_labels_file(file_path, fps=fps)
# Add metadata as attrs
ds.attrs["source_file"] = file_path.as_posix()
logger.info(f"Loaded pose tracks from {file_path}:\n{ds}")
return ds
[docs]
@register_loader("LightningPose", file_validators=[ValidDeepLabCutCSV])
def from_lp_file(file: str | Path, fps: float | None = None) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a LightningPose file.
Parameters
----------
file
Path to the file containing the predicted poses, in .csv format.
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame numbers.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
Examples
--------
>>> from movement.io import load_poses
>>> ds = load_poses.from_lp_file("path/to/file.csv", fps=30)
"""
valid_file = cast("ValidFile", file)
ds = _ds_from_lp_or_dlc_file(
valid_file=valid_file, source_software="LightningPose", fps=fps
)
n_individuals = ds.sizes.get("individual", 1)
if n_individuals > 1:
raise logger.error(
ValueError(
"LightningPose only supports single-individual datasets, "
f"but the loaded dataset has {n_individuals} individuals. "
"Did you mean to load from a DeepLabCut file instead?"
)
)
return ds
[docs]
@register_loader(
"DeepLabCut", file_validators=[ValidDeepLabCutH5, ValidDeepLabCutCSV]
)
def from_dlc_file(file: str | Path, fps: float | None = None) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a DeepLabCut file.
Parameters
----------
file
Path to the file containing the predicted poses, either in .h5
or .csv format.
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame numbers.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
See Also
--------
movement.io.load_poses.from_dlc_style_df
Examples
--------
>>> from movement.io import load_poses
>>> ds = load_poses.from_dlc_file("path/to/file.h5", fps=30)
Notes
-----
In ``movement``, pose data can only be loaded if all individuals have
the same set of keypoints (i.e., the same labeled body parts).
While DeepLabCut supports assigning keypoints that are not shared across
individuals (see the `DeepLabCut documentation for multi-animal projects
<https://deeplabcut.github.io/DeepLabCut/docs/quick-start/tutorial_maDLC.html>`_),
this feature is not currently supported in ``movement``.
"""
return _ds_from_lp_or_dlc_file(
valid_file=cast("ValidFile", file),
source_software="DeepLabCut",
fps=fps,
)
def _coco_individuals_from_categories(
frame_idx: np.ndarray,
category_ids: np.ndarray,
category_names: dict[int, str] | None,
) -> tuple[np.ndarray, list[str]]:
"""Use each COCO category as one individual across frames.
Only categories present in the results become individuals.
Returns the individual index of each detection and the individual
names, in sorted category ID order.
"""
track_ids = np.unique(category_ids)
# Index of each detection's category_id in track_ids
individual_idx = np.searchsorted(track_ids, category_ids)
# Encode each (frame, individual) slot as a single integer
# (its flat index in an n_frames x n_individuals grid) for
# checking duplicates (i.e. same individual detected multiple
# times in a frame).
pairs = frame_idx * len(track_ids) + individual_idx
if len(np.unique(pairs)) < len(pairs):
raise logger.error(
ValueError(
"COCO results contain multiple detections of the same "
"category_id in the same image_id. Set "
"category_as_track=False if categories do not "
"identify individuals."
)
)
individual_names = [
category_names[track_id]
if category_names is not None
else str(track_id)
for track_id in track_ids
]
return individual_idx, individual_names
def _coco_individuals_by_position(frame_idx: np.ndarray) -> np.ndarray:
"""Assign COCO detections to individuals by order within each frame.
Returns the individual index of each detection: the 1st detection
of a frame (in file order) is individual 0, the 2nd is individual 1,
etc.
"""
individual_idx = (
pd.Series(frame_idx).groupby(frame_idx).cumcount().to_numpy()
)
if individual_idx.max() > 0:
warnings.warn(
"COCO results do not contain cross-frame track identities. "
"Individuals are assigned positionally within each frame, "
"so identity is not guaranteed to remain stable across frames.",
stacklevel=2,
)
return individual_idx
[docs]
@register_loader("COCO", file_validators=[ValidCOCOKeypointResults])
def from_coco_file(
file: str | Path,
fps: float | None = None,
*,
annotations_file: str | Path | None = None,
category_as_track: bool = False,
use_frame_numbers_from_file: bool = False,
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a COCO keypoint results file.
Parameters
----------
file
Path to the COCO keypoint detection
`results <https://cocodataset.org/#format-results>`__ JSON file.
fps
Frames per second. If None (default), ``time`` coordinates are frame
numbers.
annotations_file
Optional path to a COCO
`annotations <https://cocodataset.org/#format-data>`__ JSON file,
used for keypoint and individual names (see Notes).
category_as_track
If True, treat each ``category_id`` as one individual tracked
across frames. If False (default), assign individuals by order of
detection within each frame (see Notes).
use_frame_numbers_from_file
Frame numbers are taken from the ``image_id`` of each detection.
If False (default), frame numbers are offset so that the lowest
image ID becomes frame 0. If True, frame numbers are kept as the
image IDs in the file. This may be useful if only a subset of the
video's frames were processed, but you want to keep the start of
the video as the time origin. In both cases, the spacing between
images is preserved, so any gaps show up in the ``time``
coordinates. For example, image IDs ``[5, 6, 9]`` in the file
become ``[0, 1, 4]`` by default and remain ``[5, 6, 9]`` if True.
Returns
-------
xarray.Dataset
``movement`` pose tracks and confidence scores.
Notes
-----
COCO keypoint results do not contain track identities. By default,
the ``i``-th detection of each image (in file order) is assigned to
individual ``i`` (named ``id_i``), so an individual does not
necessarily correspond to the same animal across frames. How to
handle this depends on your data:
- **One animal per frame**: no action needed, all detections are
assigned to ``id_0``.
- **Multiple animals, one category per animal**: set
``category_as_track=True``.
Each category present in the results then becomes one individual,
named after the category if ``annotations_file`` is provided, or
after its ``category_id`` (e.g. ``"3"``) otherwise. Multiple
detections of the same category in the same image raise a
``ValueError``.
- **Multiple animals sharing a category**: identities are not
reliable across frames. Assign them with a tracking tool before
analyses that follow individuals over time (e.g. kinematics),
and load the tracked output instead of the raw COCO results.
In ``movement``, pose data can currently only be loaded if all
individuals share the same set of keypoints. All detections in the
results must therefore have the same number of keypoints. If
``annotations_file`` is provided, all its categories must also have
the same ``keypoints`` list (same names, in the same order), with
one name per keypoint in each detection.
Otherwise, a ``ValueError`` is raised. Without an annotations file,
keypoints are assigned default names ``keypoint_0``, ``keypoint_1``,
etc. in file order.
Each keypoint in COCO results is stored as ``(x, y, v)``, where ``v``
is a visibility flag. As COCO recommends setting ``v=1`` for all
predicted keypoints, ``v`` carries no information in results files
and is discarded. Only ``x`` and ``y`` are loaded as ``position``.
The per-detection ``score`` is loaded as individual-wise
``confidence``, with dimensions ``(time, individual)``.
Examples
--------
Create a dataset from the COCO results file at "path/to/results.json",
with the time coordinates in seconds, and the keypoints named after
the categories in the annotations file at "path/to/annotations.json".
>>> from movement.io import load_poses
>>> ds = load_poses.from_coco_file(
... "path/to/results.json",
... fps=30,
... annotations_file="path/to/annotations.json",
... )
Create a dataset where each category identifies one individual tracked
across frames, with individuals named after the categories in the
annotations file.
>>> ds = load_poses.from_coco_file(
... "path/to/results.json",
... annotations_file="path/to/annotations.json",
... category_as_track=True,
... )
Create a dataset using the image IDs as frame numbers, so that
t = 0 seconds corresponds to image ID 0 (e.g. the first frame of the
full video).
>>> ds = load_poses.from_coco_file(
... "path/to/results.json",
... fps=30,
... use_frame_numbers_from_file=True,
... )
"""
valid_results = cast("ValidCOCOKeypointResults", file)
results = valid_results.data
category_names = valid_results.category_names
keypoint_names = valid_results.keypoint_names
image_ids = np.array([result["image_id"] for result in results])
category_ids = np.array([result["category_id"] for result in results])
scores = np.array([result["score"] for result in results], np.float32)
# (n_detections, n_keypoints, 3), where 3 is (x, y, visibility)
keypoints = np.array(
[result["keypoints"] for result in results], np.float32
).reshape(len(results), -1, 3)
# Locate each detection d in the output arrays by a
# (frame_idx[d], individual_idx[d]) slot.
# frame_ids: sorted unique image IDs, used as the time axis.
# frame_idx: row of each detection in frame_ids, so that
# frame_ids[frame_idx[d]] == image_ids[d].
frame_ids, frame_idx = np.unique(image_ids, return_inverse=True)
if category_as_track:
individual_idx, individual_names = _coco_individuals_from_categories(
frame_idx, category_ids, category_names
)
n_individuals = len(individual_names)
else:
individual_idx = _coco_individuals_by_position(frame_idx)
individual_names = None
n_individuals = int(individual_idx.max()) + 1
n_frames, n_keypoints = len(frame_ids), keypoints.shape[1]
position_array = np.full(
(n_frames, 2, n_keypoints, n_individuals), np.nan, np.float32
)
position_array[frame_idx, :, :, individual_idx] = keypoints[
..., :2
].transpose(0, 2, 1) # to (n_detections, 2, n_keypoints)
confidence_array = np.full((n_frames, n_individuals), np.nan, np.float32)
confidence_array[frame_idx, individual_idx] = scores
ds = from_numpy(
position_array=position_array,
confidence_array=confidence_array,
individual_names=individual_names,
keypoint_names=keypoint_names,
frame_array=(
frame_ids
if use_frame_numbers_from_file
else frame_ids - frame_ids[0]
).reshape(-1, 1),
fps=fps,
source_software="COCO",
)
ds.attrs["source_file"] = valid_results.file.as_posix()
logger.info(f"Loaded pose tracks from {valid_results.file}:\n{ds}")
return ds
def _ds_from_lp_or_dlc_file(
valid_file: ValidFile,
source_software: Literal["LightningPose", "DeepLabCut"],
fps: float | None = None,
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a LightningPose or DLC file.
Parameters
----------
valid_file
The validated LightningPose or DeepLabCut file object.
source_software
The source software of the file.
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame numbers.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
"""
# Load the DeepLabCut poses into a DataFrame
file_path = valid_file.file
df = (
_df_from_dlc_csv(valid_file)
if isinstance(valid_file, ValidDeepLabCutCSV)
else pd.DataFrame(pd.read_hdf(file_path, key="df_with_missing"))
# pd.read_hdf does not always return a DataFrame but we assume it does
# in this case (since we know what's in the "df_with_missing" dataset)
)
logger.debug(f"Loaded poses from {file_path} into a DataFrame.")
# Convert the DataFrame to an xarray dataset
ds = from_dlc_style_df(df=df, fps=fps, source_software=source_software)
# Add metadata as attrs
ds.attrs["source_file"] = file_path.as_posix()
logger.info(f"Loaded pose tracks from {file_path}:\n{ds}")
return ds
def _ds_from_sleap_analysis_file(file: Path, fps: float | None) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a SLEAP analysis (.h5) file.
Parameters
----------
file
Path to the SLEAP analysis file containing predicted pose tracks.
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame units.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
"""
with h5py.File(file, "r") as f:
# Transpose to shape: (n_frames, n_space, n_keypoints, n_tracks)
tracks = f["tracks"][:].transpose(3, 1, 2, 0)
# Create an array of NaNs for the confidence scores
scores = np.full(tracks.shape[:1] + tracks.shape[2:], np.nan)
individual_names = [n.decode() for n in f["track_names"][:]] or None
if individual_names is None:
logger.info(
f"Could not find SLEAP Track in {file}. "
"Assuming single-individual dataset and assigning "
"default individual name."
)
# If present, read the point-wise scores,
# and transpose to shape: (n_frames, n_keypoints, n_tracks)
if "point_scores" in f:
scores = f["point_scores"][:].T
return from_numpy(
position_array=tracks.astype(np.float32),
confidence_array=scores.astype(np.float32),
individual_names=individual_names,
keypoint_names=[n.decode() for n in f["node_names"][:]],
fps=fps,
source_software="SLEAP",
)
def _ds_from_sleap_labels_file(file: Path, fps: float | None) -> xr.Dataset:
"""Create a ``movement`` poses dataset from a SLEAP labels (.slp) file.
Parameters
----------
file
Path to the SLEAP labels file containing predicted pose tracks.
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame units.
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
"""
labels = read_labels(file.as_posix())
tracks_with_scores = _sleap_labels_to_numpy(labels)
individual_names = [track.name for track in labels.tracks] or None
if individual_names is None:
logger.info(
f"Could not find SLEAP Track in {file}. "
"Assuming single-individual dataset and assigning "
"default individual name."
)
return from_numpy(
position_array=tracks_with_scores[:, :-1, :, :],
confidence_array=tracks_with_scores[:, -1, :, :],
individual_names=individual_names,
keypoint_names=[kp.name for kp in labels.skeletons[0].nodes],
fps=fps,
source_software="SLEAP",
)
def _sleap_labels_to_numpy(labels: Labels) -> np.ndarray:
"""Convert a SLEAP ``Labels`` object to a NumPy array.
The output array contains pose tracks and point-wise confidence scores.
Parameters
----------
labels
A SLEAP `Labels` object.
Returns
-------
numpy.ndarray
A NumPy array containing pose tracks and confidence scores,
with shape ``(n_frames, 3, n_nodes, n_tracks)``.
Notes
-----
This function only considers SLEAP instances in the first
video of the SLEAP `Labels` object. User-labelled instances are
prioritised over predicted instances, mirroring SLEAP's approach
when exporting .h5 analysis files [1]_.
This function is adapted from `Labels.numpy()` from the
`sleap_io` package [2]_.
References
----------
.. [1] https://github.com/talmolab/sleap/blob/v1.3.3/sleap/info/write_tracking_h5.py#L59
.. [2] https://github.com/talmolab/sleap-io
"""
# Select frames from the first video only
lfs = [lf for lf in labels.labeled_frames if lf.video == labels.videos[0]]
# Figure out frame index range
frame_idxs = [lf.frame_idx for lf in lfs]
first_frame = min(0, min(frame_idxs))
last_frame = max(0, max(frame_idxs))
n_tracks = len(labels.tracks) or 1 # If no tracks, assume 1 individual
individuals = labels.tracks or [None]
skeleton = labels.skeletons[-1] # Assume project only uses last skeleton
n_nodes = len(skeleton.nodes)
n_frames = int(last_frame - first_frame + 1)
tracks = np.full((n_frames, 3, n_nodes, n_tracks), np.nan, dtype="float32")
for lf in lfs:
i = int(lf.frame_idx - first_frame)
user_instances = lf.user_instances
predicted_instances = lf.predicted_instances
for j, ind in enumerate(individuals):
user_track_instances = [
inst for inst in user_instances if inst.track == ind
]
predicted_track_instances = [
inst for inst in predicted_instances if inst.track == ind
]
# Use user-labelled instance if available
if user_track_instances:
inst = user_track_instances[-1]
tracks[i, ..., j] = np.hstack(
(inst.numpy(), np.full((n_nodes, 1), np.nan))
).T
elif predicted_track_instances:
inst = predicted_track_instances[-1]
tracks[i, ..., j] = inst.numpy(scores=True).T
return tracks
def _df_from_dlc_csv(valid_file: ValidDeepLabCutCSV) -> pd.DataFrame:
"""Create a DeepLabCut-style DataFrame from a .csv file.
If poses are loaded from a DeepLabCut-style .csv file, the DataFrame
lacks the multi-index columns that are present in the .h5 file. This
function parses the .csv file to DataFrame with multi-index columns,
i.e. the same format as in the .h5 file.
Parameters
----------
valid_file
The validated DeepLabCut-style CSV file object.
Returns
-------
pandas.DataFrame
DeepLabCut-style DataFrame with multi-index columns.
"""
# Deliberately avoid using pd.read_csv with index_col=0 here
# and instead set the index after reading the CSV,
# as in cases where the first data row is empty (e.g. "0,,,,,"),
# pandas will misinterpret that value as the index name instead of a row.
level_names = valid_file.level_names
df = pd.read_csv(
valid_file.file,
header=list(range(len(level_names))),
)
df = df.set_index(df.columns[0])
df.index.name = None
df.columns = pd.MultiIndex.from_tuples(df.columns, names=level_names) # type: ignore[arg-type]
return df
[docs]
def from_anipose_style_df(
df: pd.DataFrame,
fps: float | None = None,
individual_name: str = "id_0",
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from an Anipose 3D dataframe.
Parameters
----------
df
Anipose triangulation dataframe
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame units.
individual_name
Name of the individual, by default "id_0"
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
Notes
-----
Reshape dataframe with columns keypoint1_x, keypoint1_y, keypoint1_z,
keypoint1_score,keypoint2_x, keypoint2_y, keypoint2_z,
keypoint2_score...to array of positions with dimensions
time, space, keypoint, individual, and array of confidence (from scores)
with dimensions time, keypoint, individual.
"""
keypoint_names = sorted(
list(
set(
[
col.rsplit("_", 1)[0]
for col in df.columns
if any(col.endswith(f"_{s}") for s in ["x", "y", "z"])
]
)
)
)
n_frames = len(df)
n_keypoints = len(keypoint_names)
# Initialize arrays and fill
position_array = np.zeros(
(n_frames, 3, n_keypoints, 1)
) # 1 for single individual
confidence_array = np.zeros((n_frames, n_keypoints, 1))
for i, kp in enumerate(keypoint_names):
for j, coord in enumerate(["x", "y", "z"]):
position_array[:, j, i, 0] = df[f"{kp}_{coord}"]
confidence_array[:, i, 0] = df[f"{kp}_score"]
individual_names = [individual_name]
return from_numpy(
position_array=position_array,
confidence_array=confidence_array,
individual_names=individual_names,
keypoint_names=keypoint_names,
source_software="Anipose",
fps=fps,
)
[docs]
@register_loader("Anipose", file_validators=[ValidAniposeCSV])
def from_anipose_file(
file: str | Path,
fps: float | None = None,
individual_name: str = "id_0",
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from an Anipose 3D .csv file.
Parameters
----------
file
Path to the Anipose triangulation .csv file
fps
The number of frames per second in the video. If None (default),
the ``time`` coordinates will be in frame units.
individual_name
Name of the individual, by default "id_0"
Returns
-------
xarray.Dataset
``movement`` dataset containing the pose tracks, confidence scores,
and associated metadata.
Notes
-----
We currently do not load all information, only x, y, z, and score
(confidence) for each keypoint. Future versions will load n of cameras
and error.
"""
valid_file = cast("ValidFile", file)
anipose_df = pd.read_csv(valid_file.file)
return from_anipose_style_df(
anipose_df, fps=fps, individual_name=individual_name
)
[docs]
@register_loader("NWB", file_validators=[ValidNWBFile])
def from_nwb_file(
file: str | Path | pynwb.file.NWBFile,
processing_module_key: str = "behavior",
pose_estimation_key: str = "PoseEstimation",
) -> xr.Dataset:
"""Create a ``movement`` poses dataset from an NWB file.
The input can be a path to an NWB file on disk or a
:class:`pynwb.file.NWBFile` object.
The data will be extracted from the NWB file's
:class:`pynwb.base.ProcessingModule`
(specified by ``processing_module_key``) that contains the
``ndx_pose.PoseEstimation`` object (specified by ``pose_estimation_key``)
formatted according to the ``ndx-pose`` NWB extension [1]_.
Parameters
----------
file
Path to the NWB file on disk (ending in ".nwb"),
or an NWBFile object.
processing_module_key
Name of the :class:`pynwb.base.ProcessingModule` in the NWB file that
contains the pose estimation data. Default is "behavior".
pose_estimation_key
Name of the ``ndx_pose.PoseEstimation`` object in the processing
module (specified by ``processing_module_key``).
Default is "PoseEstimation".
Returns
-------
xarray.Dataset
A single-individual ``movement`` dataset containing the pose tracks,
confidence scores, and associated metadata.
References
----------
.. [1] https://github.com/rly/ndx-pose
Examples
--------
Open an NWB file and load pose tracks from the
:class:`pynwb.file.NWBFile` object:
>>> import pynwb
>>> import xarray as xr
>>> from movement.io import load_poses
>>> with pynwb.NWBHDF5IO("path/to/file.nwb", mode="r") as io:
... nwb_file = io.read()
... ds = load_poses.from_nwb_file(nwb_file)
Or, directly load pose tracks from an NWB file on disk:
>>> ds = load_poses.from_nwb_file("path/to/file.nwb")
Load two single-individual datasets from two NWB files and merge them
into a multi-individual dataset:
>>> ds_singles = [
... load_poses.from_nwb_file(f) for f in ["id1.nwb", "id2.nwb"]
... ]
>>> ds_multi = xr.merge(ds_singles)
"""
valid_file = cast("ValidFile", file)
file_path_or_nwbfile_obj = valid_file.file
if isinstance(file_path_or_nwbfile_obj, Path):
with pynwb.NWBHDF5IO(file_path_or_nwbfile_obj, mode="r") as io:
nwbfile_object = io.read()
ds = _ds_from_nwb_object(
nwbfile_object, processing_module_key, pose_estimation_key
)
ds.attrs["source_file"] = file_path_or_nwbfile_obj
else: # file is an NWBFile object
ds = _ds_from_nwb_object(
file_path_or_nwbfile_obj,
processing_module_key,
pose_estimation_key,
)
return ds
def _ds_from_nwb_object(
nwb_file: pynwb.file.NWBFile,
processing_module_key: str = "behavior",
pose_estimation_key: str = "PoseEstimation",
) -> xr.Dataset:
"""Extract a ``movement`` poses dataset from an NWBFile object.
Parameters
----------
nwb_file
An NWBFile object.
processing_module_key
Name of the :class:`pynwb.base.ProcessingModule` in the NWB file that
contains the pose estimation data. Default is "behavior".
pose_estimation_key
Name of the ``ndx_pose.PoseEstimation`` object in the processing
module (specified by ``processing_module_key``).
Default is "PoseEstimation".
Returns
-------
xarray.Dataset
A single-individual ``movement`` poses dataset
"""
pose_estimation = nwb_file.processing[processing_module_key][
pose_estimation_key
]
source_software = pose_estimation.source_software
pose_estimation_series = pose_estimation.pose_estimation_series
single_keypoint_datasets = []
for keypoint, pes in pose_estimation_series.items():
# Extract position and confidence data for each keypoint
position_data = np.asarray(pes.data) # shape: (n_frames, n_space)
confidence_data = ( # shape: (n_frames,)
np.asarray(pes.confidence)
if getattr(pes, "confidence", None) is not None
else np.full(position_data.shape[0], np.nan)
)
# Compute fps from time differences between timestamps
# if rate is not available
fps = pes.rate or float(np.nanmedian(1 / np.diff(pes.timestamps)))
single_keypoint_datasets.append(
# create movement dataset with 1 keypoint and 1 individual
from_numpy(
position_data[:, :, np.newaxis, np.newaxis],
confidence_data[:, np.newaxis, np.newaxis],
individual_names=[nwb_file.identifier],
keypoint_names=[keypoint],
fps=round(fps, 6),
source_software=source_software,
)
)
return xr.merge(
single_keypoint_datasets, join="outer", compat="no_conflicts"
)
