Skip to content

core.py

PoseGroup

Bases: SpyglassMixin, Manual

Groups one or more entries of keypoint pose information

Source code in src/spyglass/behavior/v1/core.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
@schema
class PoseGroup(SpyglassMixin, dj.Manual):
    """Groups one or more entries of keypoint pose information"""

    definition = """
    pose_group_name: varchar(80)
    ----
    bodyparts = NULL: longblob # list of body parts to include in the pose
    """

    class Pose(SpyglassMixinPart):
        definition = """
        -> PoseGroup
        -> PositionOutput.proj(pose_merge_id='merge_id')
        """

    def create_group(
        self,
        group_name: str,
        merge_ids: List[str],
        bodyparts: List[str] = None,
    ) -> None:
        """Create a group of pose information

        Parameters
        ----------
        group_name : str
            Name of the group
        merge_ids : List[str]
            list of keys from PositionOutput to include in the group
        bodyparts : List[str], optional
            body parts to include in the group, by default None includes all from every set
        """
        group_key = {"pose_group_name": group_name}
        if self & group_key:
            warnings.warn(
                f"Pose group {group_name} already exists. Skipping insertion."
            )
            return
        self.insert1(
            {
                **group_key,
                "bodyparts": bodyparts,
            }
        )
        for merge_id in merge_ids:
            self.Pose.insert1(
                {
                    **group_key,
                    "pose_merge_id": merge_id,
                }
            )

    def fetch_pose_datasets(
        self,
        key: dict = None,
        format_for_moseq: bool = False,
        normalize: bool = False,
        anterior_bodyparts: List[str] = None,
        posterior_bodyparts: List[str] = None,
    ):
        """fetch pose information for a group of videos

        Parameters
        ----------
        key : dict
            group key
        format_for_moseq : bool, optional
            format for MoSeq, by default False
        normalize : bool, optional
            whether to normalize the pose datasets by animal length, by default False
        anterior_bodyparts : List[str], optional
            list of anterior body parts to use for normalization, required if
            normalize is True
        posterior_bodyparts : List[str], optional
            list of posterior body parts to use for normalization, required if
            normalize is True

        Returns
        -------
        dict
            dictionary of video name to pose dataset
        """
        key = key or dict()
        self.ensure_single_entry(key)
        query = self & key
        bodyparts = query.fetch1("bodyparts")
        datasets = {}

        empty_parts = []
        if anterior_bodyparts is None or len(anterior_bodyparts) == 0:
            empty_parts.append("anterior")
        if posterior_bodyparts is None or len(posterior_bodyparts) == 0:
            empty_parts.append("posterior")
        if normalize and empty_parts:
            raise ValueError(
                f"Both anterior_bodyparts and posterior_bodyparts must be provided "
                f"as non-empty lists for normalization. Missing: {', '.join(empty_parts)}"
            )

        for merge_key in (self.Pose & query).proj(merge_id="pose_merge_id"):
            video_name = Path(
                (PositionOutput & merge_key).fetch_video_path()
            ).name
            bodyparts_df = (PositionOutput & merge_key).fetch_pose_dataframe()
            if bodyparts is None:
                bodyparts = (
                    bodyparts_df.keys().get_level_values(0).unique().values
                )
            bodyparts_df = bodyparts_df[bodyparts]
            if normalize:
                available_bodyparts = set(
                    bodyparts_df.keys().get_level_values(0).unique().values
                )
                missing_bodyparts = sorted(
                    set(anterior_bodyparts)
                    .union(posterior_bodyparts)
                    .difference(available_bodyparts)
                )
                if missing_bodyparts:
                    raise ValueError(
                        "Requested normalization bodyparts are missing from "
                        f"dataset '{video_name}': {missing_bodyparts}"
                    )
            datasets[video_name] = bodyparts_df
        if normalize:
            datasets = normalize_pose_dataset(
                datasets, anterior_bodyparts, posterior_bodyparts
            )
        if format_for_moseq:
            datasets = format_dataset_for_moseq(datasets, bodyparts)
        return datasets

    def fetch_video_paths(self, key: dict = None) -> List[Path]:
        """fetch video paths for a group of videos

        Parameters
        ----------
        key : dict
            group key

        Returns
        -------
        List[Path]
            list of video paths
        """
        key = key or dict()
        self.ensure_single_entry(key)
        key = (self & key).fetch1("KEY")
        return [
            Path((PositionOutput & merge_key).fetch_video_path())
            for merge_key in (self.Pose & key).proj(merge_id="pose_merge_id")
        ]

create_group(group_name, merge_ids, bodyparts=None)

Create a group of pose information

Parameters:

Name Type Description Default
group_name str

Name of the group

 required
merge_ids List[str]

list of keys from PositionOutput to include in the group

 required
bodyparts List[str]

body parts to include in the group, by default None includes all from every set

 None
Source code in src/spyglass/behavior/v1/core.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def create_group(
    self,
    group_name: str,
    merge_ids: List[str],
    bodyparts: List[str] = None,
) -> None:
    """Create a group of pose information

    Parameters
    ----------
    group_name : str
        Name of the group
    merge_ids : List[str]
        list of keys from PositionOutput to include in the group
    bodyparts : List[str], optional
        body parts to include in the group, by default None includes all from every set
    """
    group_key = {"pose_group_name": group_name}
    if self & group_key:
        warnings.warn(
            f"Pose group {group_name} already exists. Skipping insertion."
        )
        return
    self.insert1(
        {
            **group_key,
            "bodyparts": bodyparts,
        }
    )
    for merge_id in merge_ids:
        self.Pose.insert1(
            {
                **group_key,
                "pose_merge_id": merge_id,
            }
        )

fetch_pose_datasets(key=None, format_for_moseq=False, normalize=False, anterior_bodyparts=None, posterior_bodyparts=None)

fetch pose information for a group of videos

Parameters:

Name Type Description Default
key dict

group key

 None
format_for_moseq bool

format for MoSeq, by default False

 False
normalize bool

whether to normalize the pose datasets by animal length, by default False

 False
anterior_bodyparts List[str]

list of anterior body parts to use for normalization, required if normalize is True

 None
posterior_bodyparts List[str]

list of posterior body parts to use for normalization, required if normalize is True

 None

Returns:

Type Description
dict

dictionary of video name to pose dataset
Source code in src/spyglass/behavior/v1/core.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def fetch_pose_datasets(
    self,
    key: dict = None,
    format_for_moseq: bool = False,
    normalize: bool = False,
    anterior_bodyparts: List[str] = None,
    posterior_bodyparts: List[str] = None,
):
    """fetch pose information for a group of videos

    Parameters
    ----------
    key : dict
        group key
    format_for_moseq : bool, optional
        format for MoSeq, by default False
    normalize : bool, optional
        whether to normalize the pose datasets by animal length, by default False
    anterior_bodyparts : List[str], optional
        list of anterior body parts to use for normalization, required if
        normalize is True
    posterior_bodyparts : List[str], optional
        list of posterior body parts to use for normalization, required if
        normalize is True

    Returns
    -------
    dict
        dictionary of video name to pose dataset
    """
    key = key or dict()
    self.ensure_single_entry(key)
    query = self & key
    bodyparts = query.fetch1("bodyparts")
    datasets = {}

    empty_parts = []
    if anterior_bodyparts is None or len(anterior_bodyparts) == 0:
        empty_parts.append("anterior")
    if posterior_bodyparts is None or len(posterior_bodyparts) == 0:
        empty_parts.append("posterior")
    if normalize and empty_parts:
        raise ValueError(
            f"Both anterior_bodyparts and posterior_bodyparts must be provided "
            f"as non-empty lists for normalization. Missing: {', '.join(empty_parts)}"
        )

    for merge_key in (self.Pose & query).proj(merge_id="pose_merge_id"):
        video_name = Path(
            (PositionOutput & merge_key).fetch_video_path()
        ).name
        bodyparts_df = (PositionOutput & merge_key).fetch_pose_dataframe()
        if bodyparts is None:
            bodyparts = (
                bodyparts_df.keys().get_level_values(0).unique().values
            )
        bodyparts_df = bodyparts_df[bodyparts]
        if normalize:
            available_bodyparts = set(
                bodyparts_df.keys().get_level_values(0).unique().values
            )
            missing_bodyparts = sorted(
                set(anterior_bodyparts)
                .union(posterior_bodyparts)
                .difference(available_bodyparts)
            )
            if missing_bodyparts:
                raise ValueError(
                    "Requested normalization bodyparts are missing from "
                    f"dataset '{video_name}': {missing_bodyparts}"
                )
        datasets[video_name] = bodyparts_df
    if normalize:
        datasets = normalize_pose_dataset(
            datasets, anterior_bodyparts, posterior_bodyparts
        )
    if format_for_moseq:
        datasets = format_dataset_for_moseq(datasets, bodyparts)
    return datasets

fetch_video_paths(key=None)

fetch video paths for a group of videos

Parameters:

Name Type Description Default
key dict

group key

 None

Returns:

Type Description
List[Path]

list of video paths
Source code in src/spyglass/behavior/v1/core.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def fetch_video_paths(self, key: dict = None) -> List[Path]:
    """fetch video paths for a group of videos

    Parameters
    ----------
    key : dict
        group key

    Returns
    -------
    List[Path]
        list of video paths
    """
    key = key or dict()
    self.ensure_single_entry(key)
    key = (self & key).fetch1("KEY")
    return [
        Path((PositionOutput & merge_key).fetch_video_path())
        for merge_key in (self.Pose & key).proj(merge_id="pose_merge_id")
    ]

normalize_pose_dataset(datasets, anterior_bodyparts, posterior_bodyparts)

Normalize pose datasets by centering and scaling based on the mean position and length of the anterior and posterior body parts

Parameters:

Name Type Description Default
datasets dict[str, DataFrame]

dictionary of video name to pose dataset (as returned by fetch_pose_datasets)

 required
anterior_bodyparts List[str]

list of anterior body parts to use for normalization

 required
posterior_bodyparts List[str]

list of posterior body parts to use for normalization

 required

Returns:

Type Description
dict[str, DataFrame]

normalized pose datasets
Source code in src/spyglass/behavior/v1/core.py 
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
def normalize_pose_dataset(
    datasets: dict[str, pd.DataFrame],
    anterior_bodyparts: List[str],
    posterior_bodyparts: List[str],
) -> dict[str, pd.DataFrame]:
    """
    Normalize pose datasets by centering and scaling based on the mean position and length
    of the anterior and posterior body parts

    Parameters
    ----------
    datasets : dict[str, pd.DataFrame]
        dictionary of video name to pose dataset (as returned by fetch_pose_datasets)
    anterior_bodyparts : List[str]
        list of anterior body parts to use for normalization
    posterior_bodyparts : List[str]
        list of posterior body parts to use for normalization

    Returns
    -------
    dict[str, pd.DataFrame]
        normalized pose datasets
    """

    for video, dataset in datasets.items():
        datasets[video] = _normalize_1_pose_dataset(
            dataset, anterior_bodyparts, posterior_bodyparts
        )
    return datasets

format_dataset_for_moseq(datasets, bodyparts, coordinate_axes=['x', 'y'])

format pose datasets for MoSeq

Parameters:

Name Type Description Default
datasets dict[str, DataFrame]

dictionary of video name to pose dataset

 required
bodyparts List[str]

list of body parts to include in the pose

 required
coordinate_axes List[str]

list of coordinate axes to include, by default ["x", "y"]

 ['x', 'y']

Returns:

Type Description
tuple[dict[str, np.ndarray], dict[str, np.ndarray]

coordinates and confidences for each video
Source code in src/spyglass/behavior/v1/core.py 
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
def format_dataset_for_moseq(
    datasets: dict[str, pd.DataFrame],
    bodyparts: List[str],
    coordinate_axes: List[str] = ["x", "y"],
):
    """format pose datasets for MoSeq

    Parameters
    ----------
    datasets : dict[str, pd.DataFrame]
        dictionary of video name to pose dataset
    bodyparts : List[str]
        list of body parts to include in the pose
    coordinate_axes : List[str], optional
        list of coordinate axes to include, by default ["x", "y"]

    Returns
    -------
    tuple[dict[str, np.ndarray], dict[str, np.ndarray]
        coordinates and confidences for each video
    """
    num_keypoints = len(bodyparts)
    num_dimensions = len(coordinate_axes)
    coordinates = {}
    confidences = {}

    for video_name, bodyparts_df in datasets.items():
        num_frames = len(bodyparts_df[bodyparts[0]])
        coordinates_i = np.empty((num_frames, num_keypoints, num_dimensions))
        confidences_i = np.empty((num_frames, num_keypoints))

        for i, bodypart in enumerate(bodyparts):
            part_df = bodyparts_df[bodypart]
            coordinates_i[:, i, :] = part_df[coordinate_axes].values
            confidences_i[:, i] = part_df["likelihood"].values
        coordinates[video_name] = coordinates_i
        confidences[video_name] = confidences_i
    return coordinates, confidences