Social rats - PAIR-R24M#

Data from Marshall et al. [2021] (figshare), a multi-animal 3D pose dataset about the dyadic interactions in laboratory rats.

../_images/pair24.png

Figure taken from Fig. 1 of Marshall et al. [2021]

from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd
import xarray as xr
from movement.kinematics import compute_speed, compute_velocity
from movement.utils.vector import compute_norm

import ethograph as eto
from ethograph.io.pairing import pair_media

try:
    _here = Path(__vsc_ipynb_file__).parent
except NameError:
    _here = Path().resolve()

SESSION = "20210119_Recording_SR1_SR2_social_vidtwo"


# The project folder holds what spans sessions (mapping.txt, pipeline configs);
# each session folder sits inside it. Choose PROJECT_DIR on the GUI's start page.
PROJECT_DIR = _here.parent / "data" / "pair24"
DATA_DIR = PROJECT_DIR / SESSION
VIDEO_DIR = DATA_DIR / "videos"

# TODO: adjust to your local copy of the PAIR-R24M dataset
path = DATA_DIR / "markerDataset.csv"

CAMERAS = ["Camera1", "Camera2", "Camera3", "Camera4", "Camera5", "Camera6"]
FPS = 120

def from_pair24_csv(
    file_path: Path | str,
    fps: float,
) -> xr.Dataset:
    """The whole recording as one dataset on the session clock (frame i at i / fps)."""
    df = pd.read_csv(file_path)

    keypoint_names = [
        "HeadF",
        "HeadB",
        "HeadL",
        "SpineF",
        "SpineM",
        "SpineL",
        "Offset1",
        "Offset2",
        "HipL",
        "HipR",
        "ShoulderL",
        "ShoulderR",
    ]
    individual_names = ["an1", "an2"]
    position_types = ["aligned", "absolute"]
    n_frames = len(df)
    n_keypoints = len(keypoint_names)
    n_individuals = len(individual_names)
    n_space = 3

    position_array = np.zeros((n_frames, len(position_types), n_space, n_keypoints, n_individuals))
    confidence_array = np.ones((n_frames, len(position_types), n_keypoints, n_individuals))

    for p, pos_type in enumerate(position_types):
        csv_prefix = "alignedPosition" if pos_type == "aligned" else "absolutePosition"
        for i, individual in enumerate(individual_names):
            for j, keypoint in enumerate(keypoint_names):
                for k, coord in enumerate(["x", "y", "z"]):
                    col_name = f"{csv_prefix}_{individual}_{keypoint}_{coord}"
                    if col_name in df.columns:
                        position_array[:, p, k, j, i] = df[col_name].values

    time_coords = np.arange(n_frames, dtype=float) / fps

    ds = xr.Dataset(
        data_vars={
            "position": xr.DataArray(
                position_array,
                dims=["time", "position_type", "space", "keypoints", "individual"],
            ),
            "confidence": xr.DataArray(
                confidence_array,
                dims=["time", "position_type", "keypoints", "individual"],
            ),
        },
        coords={
            "time": time_coords,
            "position_type": position_types,
            "space": ["x", "y", "z"],
            "keypoints": keypoint_names,
            "individual": ["mouse 1", "mouse 2"],
        },
        attrs={"source_software": "DeepLabCut", "fps": fps},
    )

    com_data = np.zeros((n_frames, n_space, n_individuals))
    for i, individual in enumerate(individual_names):
        for j, coord in enumerate(["x", "y", "z"]):
            col_name = f"centerOfmass_{individual}_{coord}"
            if col_name in df.columns:
                com_data[:, j, i] = df[col_name].values
    ds["center_of_mass"] = xr.DataArray(com_data, dims=["time", "space", "individual"])

    ds["pairwise_distance"] = compute_norm(
        ds.center_of_mass.sel(individual="mouse 1") - ds.center_of_mass.sel(individual="mouse 2")
    )
    ds["nose_nose_distance"] = compute_norm(
        ds.position.sel(keypoints="HeadF", individual="mouse 1", position_type="absolute")
        - ds.position.sel(keypoints="HeadF", individual="mouse 2", position_type="absolute")
    )
    ds["velocity"] = compute_velocity(ds.position.sel(position_type="aligned"))
    ds["speed"] = compute_speed(ds.position.sel(position_type="aligned"))

    return ds


def trial_table_from_videos(
    video_dir: Path,
    cameras: list[str],
    n_frames: int,
    fps: float,
) -> pd.DataFrame:
    """One trial per video chunk; each file is named by its first frame on the session clock.

    A chunk that starts after the mocap's last frame has no data and is dropped, and the
    last trial ends where the mocap does. The timing goes into the alignment, which is
    the one source of trial timing: the GUI slices the long dataset by it.
    """
    starts = sorted(int(f.stem) for f in (video_dir / cameras[0]).glob("*.mp4"))
    starts = [s for s in starts if s < n_frames]
    stops = [min(s, n_frames) for s in starts[1:]] + [n_frames]

    rows = []
    for trial, (start, stop) in enumerate(zip(starts, stops), start=1):
        # Full path: pair_media stores the basename per trial and the path per stream.
        row = {"trial": trial, "start_time": start / fps, "stop_time": stop / fps}
        for cam in cameras:
            row[f"video_{cam}"] = str(video_dir / cam / f"{start}.mp4")
        rows.append(row)
    return pd.DataFrame(rows)

ds = from_pair24_csv(path, fps=FPS)
ds

<xarray.Dataset> Size: 257MB
Dimensions:             (time: 108000, position_type: 2, space: 3,
                         keypoints: 12, individual: 2)
Coordinates:
  * time                (time) float64 864kB 0.0 0.008333 ... 900.0 900.0
  * position_type       (position_type) <U8 64B 'aligned' 'absolute'
  * space               (space) <U1 12B 'x' 'y' 'z'
  * keypoints           (keypoints) <U9 432B 'HeadF' 'HeadB' ... 'ShoulderR'
  * individual          (individual) <U7 56B 'mouse 1' 'mouse 2'
Data variables:
    position            (time, position_type, space, keypoints, individual) float64 124MB ...
    confidence          (time, position_type, keypoints, individual) float64 41MB ...
    center_of_mass      (time, space, individual) float64 5MB -183.0 ... 83.7
    pairwise_distance   (time) float64 864kB 56.41 54.29 52.29 ... 273.1 273.1
    nose_nose_distance  (time) float64 864kB 110.6 110.3 109.4 ... 129.3 129.2
    velocity            (time, space, keypoints, individual) float64 62MB 1.5...
    speed               (time, keypoints, individual) float64 21MB 2.017e+03 ...
Attributes:
    source_software:  DeepLabCut
    fps:              120
session_table = trial_table_from_videos(VIDEO_DIR, CAMERAS, n_frames=ds.sizes["time"], fps=FPS)
session_table
trial start_time stop_time video_Camera1 video_Camera2 video_Camera3 video_Camera4 video_Camera5 video_Camera6
0 1 0.000000 29.166667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
1 2 29.166667 58.333333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
2 3 58.333333 87.500000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
3 4 87.500000 116.666667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
4 5 116.666667 145.833333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
5 6 145.833333 175.000000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
6 7 175.000000 204.166667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
7 8 204.166667 233.333333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
8 9 233.333333 262.500000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
9 10 262.500000 291.666667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
10 11 291.666667 320.833333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
11 12 320.833333 350.000000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
12 13 350.000000 379.166667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
13 14 379.166667 408.333333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
14 15 408.333333 437.500000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
15 16 437.500000 466.666667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
16 17 466.666667 495.833333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
17 18 495.833333 525.000000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
18 19 525.000000 554.166667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
19 20 554.166667 583.333333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
20 21 583.333333 612.500000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
21 22 612.500000 641.666667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
22 23 641.666667 670.833333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
23 24 670.833333 700.000000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
24 25 700.000000 729.166667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
25 26 729.166667 758.333333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
26 27 758.333333 787.500000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
27 28 787.500000 816.666667 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
28 29 816.666667 845.833333 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
29 30 845.833333 875.000000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p...
30 31 875.000000 900.000000 c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... c:\Users\aksel\Documents\Code\ethograph\data\p... 
output_path = DATA_DIR / "Trial_data.nc"
output_path.parent.mkdir(parents=True, exist_ok=True)

# The alignment holds the trials (start/stop on the session clock) and the videos per trial.
nwb_path = DATA_DIR / ".ethograph" / "alignment.nwb"
pair_media(
    trial_table=session_table,
    stream_rates={"video": float(FPS)},
    output_path=nwb_path,
    individuals=[str(v) for v in ds.individual.values],
    on_existing="replace",
)

# The dataset stays one long recording; no per-trial split.
ds.to_netcdf(output_path)
print(f"Saved dataset to {output_path}")

Saved dataset to c:\Users\aksel\Documents\Code\ethograph\data\pair24\20210119_Recording_SR1_SR2_social_vidtwo\Trial_data.nc

from ethograph.labels.converters import write_mapping_file

mapping = {
    "Background": 0,
    "Idle": 1,
    "SmallMovement": 2,
    "HeadTilt": 3,
    "Groom": 4,
    "Sniff": 5,
    "Investigate": 6,
    "RearUp": 7,
    "RearDown": 8,
    "CrouchExplore": 9,
    "Amble": 10,
    "Locomotion": 11,
}

# The label vocabulary belongs to the project, shared by every session in it.
# A {session}/.ethograph/mapping.txt would override it for that session only.
mapping_path = PROJECT_DIR / / ".ethograph" / "mapping.txt"
write_mapping_file(mapping_path, mapping)
print(f"Saved mapping to {mapping_path}")

Saved mapping to c:\Users\aksel\Documents\Code\ethograph\data\pair24\mapping.txt