Skip to content

deeplabcut.pose_estimation_pytorch.runners.shelving

Modules used to read/write shelve data during video analysis in DeepLabCut 3.0.

Classes:

Name Description
FeatureShelfWriter

Writes bodypart features to a shelf on-the-fly for ReID model training.
ShelfManager

Class to manage shelf data.
ShelfReader

Reads data from a shelf.
ShelfWriter

Writes data to a shelf on-the-fly during video analysis.

FeatureShelfWriter

Bases: ShelfWriter

Writes bodypart features to a shelf on-the-fly for ReID model training.

Parameters:

Name Type Description Default

pose_cfg

dict

The test pose config for the model.

 required

filepath

str | Path

The path where the data should be saved.

 required

num_frames

int | None

The number of frames in the video. Used to set the number of leading 0s in the keys of the dictionary. Default is 5 if the number of frames is not given.

 None

Attributes:

Name Type Description
filepath

The path to the shelf.

Methods:

Name Description
add_prediction

Adds the prediction for a frame to the shelf.
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
class FeatureShelfWriter(ShelfWriter):
    """Writes bodypart features to a shelf on-the-fly for ReID model training.

    Args:
        pose_cfg: The test pose config for the model.
        filepath: The path where the data should be saved.
        num_frames: The number of frames in the video. Used to set the number of
            leading 0s in the keys of the dictionary. Default is 5 if the number of
            frames is not given.

    Attributes:
        filepath: The path to the shelf.
    """

    def __init__(self, pose_cfg: dict, filepath: str | Path, num_frames: int | None = None):
        super().__init__(pose_cfg, filepath, num_frames)

    def add_prediction(
        self,
        bodyparts: np.ndarray,
        features: np.ndarray | None = None,
        **kwargs,
    ) -> None:
        """Adds the prediction for a frame to the shelf.

        Args:
            bodyparts: The predicted bodyparts.
            features: The features for the bodyparts.
        """
        if not self._open:
            raise ValueError("You must call open() before adding data!")

        key = "frame" + str(self._frame_index).zfill(self._str_width)

        # bodyparts to shape (num_assemblies, num_bpts, xy)
        coordinates = bodyparts[:, :, :2]
        if features is None:
            raise ValueError("Backbone features must be given to the FeatureShelfWriter")

        self._db[key] = dict(coordinates=coordinates, features=features)
        self._frame_index += 1

add_prediction 

add_prediction(bodyparts: ndarray, features: ndarray | None = None, **kwargs) -> None

Adds the prediction for a frame to the shelf.

Parameters:

Name Type Description Default

bodyparts

ndarray

The predicted bodyparts.

 required

features

ndarray | None

The features for the bodyparts.

 None
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def add_prediction(
    self,
    bodyparts: np.ndarray,
    features: np.ndarray | None = None,
    **kwargs,
) -> None:
    """Adds the prediction for a frame to the shelf.

    Args:
        bodyparts: The predicted bodyparts.
        features: The features for the bodyparts.
    """
    if not self._open:
        raise ValueError("You must call open() before adding data!")

    key = "frame" + str(self._frame_index).zfill(self._str_width)

    # bodyparts to shape (num_assemblies, num_bpts, xy)
    coordinates = bodyparts[:, :, :2]
    if features is None:
        raise ValueError("Backbone features must be given to the FeatureShelfWriter")

    self._db[key] = dict(coordinates=coordinates, features=features)
    self._frame_index += 1

ShelfManager

Class to manage shelf data.

Methods:

Name Description
close

Closes the shelf.
open

Opens the shelf.
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
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
class ShelfManager:
    """Class to manage shelf data."""

    def __init__(self, filepath: str | Path, flag: str = "r") -> None:
        self.filepath = Path(filepath)
        self.flag = flag

        self._db: shelve.Shelf | None = None
        self._open: bool = False

    def open(self) -> None:
        """Opens the shelf."""
        self._db = shelve.open(
            str(self.filepath),
            flag=self.flag,
            protocol=pickle.DEFAULT_PROTOCOL,
        )
        self._open = True

    def close(self) -> None:
        """Closes the shelf."""
        if not self._open:
            return

        try:
            self._db.close()
        except AttributeError:
            pass

        self._open = False

    def keys(self) -> list[str]:
        if not self._open:
            raise ValueError("You must call open() before reading keys!")

        return [k for k in self._db]

close 

close() -> None

Closes the shelf.

Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
39
40
41
42
43
44
45
46
47
48
49
def close(self) -> None:
    """Closes the shelf."""
    if not self._open:
        return

    try:
        self._db.close()
    except AttributeError:
        pass

    self._open = False

open 

open() -> None

Opens the shelf.

Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
30
31
32
33
34
35
36
37
def open(self) -> None:
    """Opens the shelf."""
    self._db = shelve.open(
        str(self.filepath),
        flag=self.flag,
        protocol=pickle.DEFAULT_PROTOCOL,
    )
    self._open = True

ShelfReader

Bases: ShelfManager

Reads data from a shelf.

Methods:

Name Description
__getitem__

Reads an item from the shelf.
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class ShelfReader(ShelfManager):
    """Reads data from a shelf."""

    def __getitem__(self, item: str) -> dict:
        """Reads an item from the shelf.

        Args:
            item: The key of the item to read.

        Returns:
            The item.
        """
        if not self._open:
            raise ValueError("You must call open() before reading data!")

        return self._db[item]

__getitem__ 

__getitem__(item: str) -> dict

Reads an item from the shelf.

Parameters:

Name Type Description Default

item

str

The key of the item to read.

 required

Returns:

Type Description
dict

The item.
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
def __getitem__(self, item: str) -> dict:
    """Reads an item from the shelf.

    Args:
        item: The key of the item to read.

    Returns:
        The item.
    """
    if not self._open:
        raise ValueError("You must call open() before reading data!")

    return self._db[item]

ShelfWriter

Bases: ShelfManager

Writes data to a shelf on-the-fly during video analysis.

Parameters:

Name Type Description Default

pose_cfg

dict

The test pose config for the model.

 required

filepath

str | Path

The path where the data should be saved.

 required

num_frames

int | None

The number of frames in the video. Used to set the number of leading 0s in the keys of the dictionary. Default is 5 if the number of frames is not given.

 None

Attributes:

Name Type Description
filepath

The path to the shelf.

Methods:

Name Description
add_prediction

Adds the prediction for a frame to the shelf.
close

Closes the shelf and writes the updated metadata.
open

Opens the shelf.
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
 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
168
169
170
171
172
173
174
175
class ShelfWriter(ShelfManager):
    """Writes data to a shelf on-the-fly during video analysis.

    Args:
        pose_cfg: The test pose config for the model.
        filepath: The path where the data should be saved.
        num_frames: The number of frames in the video. Used to set the number of
            leading 0s in the keys of the dictionary. Default is 5 if the number of
            frames is not given.

    Attributes:
        filepath: The path to the shelf.
    """

    def __init__(self, pose_cfg: dict, filepath: str | Path, num_frames: int | None = None):
        super().__init__(filepath, flag="c")
        self._pose_cfg = pose_cfg
        self._num_frames = num_frames
        self._frame_index = 0

        self._str_width = 5
        if num_frames is not None:
            self._str_width = int(np.ceil(np.log10(num_frames)))

    def add_prediction(
        self,
        bodyparts: np.ndarray,
        unique_bodyparts: np.ndarray | None = None,
        identity_scores: np.ndarray | None = None,
        **kwargs,
    ) -> None:
        """Adds the prediction for a frame to the shelf.

        Args:
            bodyparts: The predicted bodyparts.
            unique_bodyparts: The predicted unique bodyparts, if there are any.
            identity_scores: The predicted identities, if there are any.
        """
        if not self._open:
            raise ValueError("You must call open() before adding data!")

        key = "frame" + str(self._frame_index).zfill(self._str_width)

        # convert bodyparts to shape (num_bpts, num_assemblies, 3)
        bodyparts = bodyparts.transpose((1, 0, 2))
        coordinates = [bpt[:, :2] for bpt in bodyparts]
        scores = [bpt[:, 2:3] for bpt in bodyparts]

        # full pickle has bodyparts and unique bodyparts in same array
        if unique_bodyparts is not None:
            unique_bpts = unique_bodyparts.transpose((1, 0, 2))
            coordinates += [bpt[:, :2] for bpt in unique_bpts]
            scores += [bpt[:, 2:] for bpt in unique_bpts]

        output = dict(coordinates=(coordinates,), confidence=scores, costs=None)

        if identity_scores is not None:
            # Reshape id scores from (num_assemblies, num_bpts, num_individuals)
            # to the original DLC full pickle format: (num_bpts, num_assem, num_ind)
            id_scores = identity_scores.transpose((1, 0, 2))
            output["identity"] = [bpt_id_scores for bpt_id_scores in id_scores]

            if unique_bodyparts is not None:
                # needed for create_video_with_all_detections to display unique bpts
                num_unique = unique_bodyparts.shape[1]
                num_assem, num_ind = id_scores.shape[1:]
                output["identity"] += [-1 * np.ones((num_assem, num_ind)) for i in range(num_unique)]

        self._db[key] = output
        self._frame_index += 1

    def close(self) -> None:
        """Closes the shelf and writes the updated metadata."""
        if self._open and self._frame_index > 0:
            # Write updated metadata to shelf (top-level indexing required for shelve)
            metadata = self._db["metadata"]
            metadata["nframes"] = self._frame_index
            self._db["metadata"] = metadata

        super().close()

    def open(self) -> None:
        """Opens the shelf."""
        super().open()
        self._frame_index = 0

        all_joints = self._pose_cfg["all_joints"]
        paf_graph = self._pose_cfg.get("partaffinityfield_graph", [])

        self._db["metadata"] = {
            "nms radius": self._pose_cfg.get("nmsradius"),
            "minimal confidence": self._pose_cfg.get("minconfidence"),
            "sigma": self._pose_cfg.get("sigma", 1),
            "PAFgraph": paf_graph,
            "PAFinds": self._pose_cfg.get("paf_best", np.arange(len(paf_graph))),
            "all_joints": [[i] for i in range(len(all_joints))],
            "all_joints_names": [self._pose_cfg["all_joints_names"][i] for i in range(len(all_joints))],
            "nframes": self._num_frames,
            "key_str_width": self._str_width,
        }

add_prediction 

add_prediction(
    bodyparts: ndarray, unique_bodyparts: ndarray | None = None, identity_scores: ndarray | None = None, **kwargs
) -> None

Adds the prediction for a frame to the shelf.

Parameters:

Name Type Description Default

bodyparts

ndarray

The predicted bodyparts.

 required

unique_bodyparts

ndarray | None

The predicted unique bodyparts, if there are any.

 None

identity_scores

ndarray | None

The predicted identities, if there are any.

 None
Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
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
def add_prediction(
    self,
    bodyparts: np.ndarray,
    unique_bodyparts: np.ndarray | None = None,
    identity_scores: np.ndarray | None = None,
    **kwargs,
) -> None:
    """Adds the prediction for a frame to the shelf.

    Args:
        bodyparts: The predicted bodyparts.
        unique_bodyparts: The predicted unique bodyparts, if there are any.
        identity_scores: The predicted identities, if there are any.
    """
    if not self._open:
        raise ValueError("You must call open() before adding data!")

    key = "frame" + str(self._frame_index).zfill(self._str_width)

    # convert bodyparts to shape (num_bpts, num_assemblies, 3)
    bodyparts = bodyparts.transpose((1, 0, 2))
    coordinates = [bpt[:, :2] for bpt in bodyparts]
    scores = [bpt[:, 2:3] for bpt in bodyparts]

    # full pickle has bodyparts and unique bodyparts in same array
    if unique_bodyparts is not None:
        unique_bpts = unique_bodyparts.transpose((1, 0, 2))
        coordinates += [bpt[:, :2] for bpt in unique_bpts]
        scores += [bpt[:, 2:] for bpt in unique_bpts]

    output = dict(coordinates=(coordinates,), confidence=scores, costs=None)

    if identity_scores is not None:
        # Reshape id scores from (num_assemblies, num_bpts, num_individuals)
        # to the original DLC full pickle format: (num_bpts, num_assem, num_ind)
        id_scores = identity_scores.transpose((1, 0, 2))
        output["identity"] = [bpt_id_scores for bpt_id_scores in id_scores]

        if unique_bodyparts is not None:
            # needed for create_video_with_all_detections to display unique bpts
            num_unique = unique_bodyparts.shape[1]
            num_assem, num_ind = id_scores.shape[1:]
            output["identity"] += [-1 * np.ones((num_assem, num_ind)) for i in range(num_unique)]

    self._db[key] = output
    self._frame_index += 1

close 

close() -> None

Closes the shelf and writes the updated metadata.

Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
147
148
149
150
151
152
153
154
155
def close(self) -> None:
    """Closes the shelf and writes the updated metadata."""
    if self._open and self._frame_index > 0:
        # Write updated metadata to shelf (top-level indexing required for shelve)
        metadata = self._db["metadata"]
        metadata["nframes"] = self._frame_index
        self._db["metadata"] = metadata

    super().close()

open 

open() -> None

Opens the shelf.

Source code in deeplabcut/pose_estimation_pytorch/runners/shelving.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def open(self) -> None:
    """Opens the shelf."""
    super().open()
    self._frame_index = 0

    all_joints = self._pose_cfg["all_joints"]
    paf_graph = self._pose_cfg.get("partaffinityfield_graph", [])

    self._db["metadata"] = {
        "nms radius": self._pose_cfg.get("nmsradius"),
        "minimal confidence": self._pose_cfg.get("minconfidence"),
        "sigma": self._pose_cfg.get("sigma", 1),
        "PAFgraph": paf_graph,
        "PAFinds": self._pose_cfg.get("paf_best", np.arange(len(paf_graph))),
        "all_joints": [[i] for i in range(len(all_joints))],
        "all_joints_names": [self._pose_cfg["all_joints_names"][i] for i in range(len(all_joints))],
        "nframes": self._num_frames,
        "key_str_width": self._str_width,
    }