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
@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
        keys : List[dict]
            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
    ):
        """fetch pose information for a group of videos

        Parameters
        ----------
        key : dict
            group key
        format_for_moseq : bool, optional
            format for MoSeq, by default False

        Returns
        -------
        dict
            dictionary of video name to pose dataset
        """
        if key is None:
            key = {}

        bodyparts = (self & key).fetch1("bodyparts")
        datasets = {}
        for merge_key in (self.Pose & key).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]
            datasets[video_name] = bodyparts_df
        if format_for_moseq:
            datasets = format_dataset_for_moseq(datasets, bodyparts)
        return datasets

    def fetch_video_paths(self, key: dict = None):
        """fetch video paths for a group of videos

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

        Returns
        -------
        List[Path]
            list of video paths
        """
        if key is None:
            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
keys List[dict]

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
    keys : List[dict]
        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)

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

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
def fetch_pose_datasets(
    self, key: dict = None, format_for_moseq: bool = False
):
    """fetch pose information for a group of videos

    Parameters
    ----------
    key : dict
        group key
    format_for_moseq : bool, optional
        format for MoSeq, by default False

    Returns
    -------
    dict
        dictionary of video name to pose dataset
    """
    if key is None:
        key = {}

    bodyparts = (self & key).fetch1("bodyparts")
    datasets = {}
    for merge_key in (self.Pose & key).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]
        datasets[video_name] = bodyparts_df
    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 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def fetch_video_paths(self, key: dict = None):
    """fetch video paths for a group of videos

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

    Returns
    -------
    List[Path]
        list of video paths
    """
    if key is None:
        key = {}
    return [
        Path((PositionOutput & merge_key).fetch_video_path())
        for merge_key in (self.Pose & key).proj(merge_id="pose_merge_id")
    ]

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 
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
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