Skip to content

deeplabcut.pose_estimation_pytorch.models.modules.kpt_encoders

Classes:

Name Description
BaseKeypointEncoder

Encodes keypoints into heatmaps.
ColoredKeypointEncoder

Encodes keypoints into a given number of color channels.
StackedKeypointEncoder

Encodes keypoints into heatmaps, where each.

BaseKeypointEncoder

Bases: ABC

Encodes keypoints into heatmaps.

Modified from BUCTD/data/JointsDataset

Methods:

Name Description
__call__

Args:
__init__

Args:
blur_heatmap

Applies a Gaussian blur to a heatmap.
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
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
class BaseKeypointEncoder(ABC):
    """Encodes keypoints into heatmaps.

    Modified from BUCTD/data/JointsDataset
    """

    def __init__(
        self,
        num_joints: int,
        kernel_size: tuple[int, int] = (15, 15),
        img_size: tuple[int, int] = (256, 256),
    ) -> None:
        """
        Args:
            num_joints: The number of joints to encode
            kernel_size: The Gaussian kernel size to use when blurring a heatmap
            img_size: The (height, width) of the input images
        """
        self.kernel_size = kernel_size
        self.num_joints = num_joints
        self.img_size = img_size

    @property
    @abstractmethod
    def num_channels(self):
        pass

    @abstractmethod
    def __call__(self, keypoints: np.ndarray, size: tuple[int, int]) -> np.ndarray:
        """
        Args:
            keypoints: the keypoints to encode
            size: the (height, width) of the heatmap in which the keypoints should
                be encoded

        Returns:
            the encoded keypoints
        """
        raise NotImplementedError

    def blur_heatmap(self, heatmap: np.ndarray) -> np.ndarray:
        """Applies a Gaussian blur to a heatmap.

        Taken from BUCTD/data/JointsDataset, generate_heatmap

        Args:
            heatmap: the heatmap to blur (with values in [0, 1] or [0, 255])

        Returns:
            The heatmap with a Gaussian blur, such that max(heatmap) = 255
        """
        heatmap = cv2.GaussianBlur(heatmap, self.kernel_size, sigmaX=0)
        am = np.amax(heatmap)
        if am == 0:
            return heatmap
        heatmap /= am / 255
        return heatmap

__call__ abstractmethod 

__call__(keypoints: ndarray, size: tuple[int, int]) -> np.ndarray

Parameters:

Name Type Description Default

keypoints

ndarray

the keypoints to encode

 required

size

tuple[int, int]

the (height, width) of the heatmap in which the keypoints should be encoded

 required

Returns:

Type Description
ndarray

the encoded keypoints
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
52
53
54
55
56
57
58
59
60
61
62
63
@abstractmethod
def __call__(self, keypoints: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """
    Args:
        keypoints: the keypoints to encode
        size: the (height, width) of the heatmap in which the keypoints should
            be encoded

    Returns:
        the encoded keypoints
    """
    raise NotImplementedError

__init__ 

__init__(num_joints: int, kernel_size: tuple[int, int] = (15, 15), img_size: tuple[int, int] = (256, 256)) -> None

Parameters:

Name Type Description Default

num_joints

int

The number of joints to encode

 required

kernel_size

tuple[int, int]

The Gaussian kernel size to use when blurring a heatmap

 (15, 15)

img_size

tuple[int, int]

The (height, width) of the input images

 (256, 256)
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(
    self,
    num_joints: int,
    kernel_size: tuple[int, int] = (15, 15),
    img_size: tuple[int, int] = (256, 256),
) -> None:
    """
    Args:
        num_joints: The number of joints to encode
        kernel_size: The Gaussian kernel size to use when blurring a heatmap
        img_size: The (height, width) of the input images
    """
    self.kernel_size = kernel_size
    self.num_joints = num_joints
    self.img_size = img_size

blur_heatmap 

blur_heatmap(heatmap: ndarray) -> np.ndarray

Applies a Gaussian blur to a heatmap.

Taken from BUCTD/data/JointsDataset, generate_heatmap

Parameters:

Name Type Description Default

heatmap

ndarray

the heatmap to blur (with values in [0, 1] or [0, 255])

 required

Returns:

Type Description
ndarray

The heatmap with a Gaussian blur, such that max(heatmap) = 255
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def blur_heatmap(self, heatmap: np.ndarray) -> np.ndarray:
    """Applies a Gaussian blur to a heatmap.

    Taken from BUCTD/data/JointsDataset, generate_heatmap

    Args:
        heatmap: the heatmap to blur (with values in [0, 1] or [0, 255])

    Returns:
        The heatmap with a Gaussian blur, such that max(heatmap) = 255
    """
    heatmap = cv2.GaussianBlur(heatmap, self.kernel_size, sigmaX=0)
    am = np.amax(heatmap)
    if am == 0:
        return heatmap
    heatmap /= am / 255
    return heatmap

ColoredKeypointEncoder

Bases: BaseKeypointEncoder

Encodes keypoints into a given number of color channels.

Modified from BUCTD/data/JointsDataset, get_condition_image_colored

Methods:

Name Description
__call__

Args:
__init__

Args:
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
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
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
@KEYPOINT_ENCODERS.register_module
class ColoredKeypointEncoder(BaseKeypointEncoder):
    """Encodes keypoints into a given number of color channels.

    Modified from BUCTD/data/JointsDataset, get_condition_image_colored
    """

    def __init__(self, colors: list[tuple[int, int, int]] | None = None, **kwargs) -> None:
        """
        Args:
            colors: the color to use for each keypoint
        """
        super().__init__(**kwargs)
        if colors is None:
            colors = self.get_colors_from_cmap("rainbow", self.num_joints)
        self.colors = np.array(colors)

    @property
    def num_channels(self):
        return 3

    def __call__(self, keypoints: np.ndarray, size: tuple[int, int]) -> np.ndarray:
        """
        Args:
            keypoints: batch of keypoints to encode with shape (batch_size, num_joints, 2)
            size: the (height, width) of the heatmap in which the keypoints should be encoded

        Returns:
            encoded keypoints with shape (batch_size, num_joints, height, width, 3)
        """

        batch_size, num_kpts, _ = keypoints.shape

        if not num_kpts == len(self.colors):
            raise ValueError(
                f"Cannot encode the keypoints. Initialized with {len(self.colors)} "
                f"colors, but there are {num_kpts} to encode"
            )

        # kpts = keypoints.detach().numpy()
        kpts = keypoints.copy()
        kpts[keypoints[..., 2] <= 0] = 0

        # Mark keypoints as visible, remove NaNs
        kpts[kpts[..., 2] > 0, 2] = 2
        kpts = np.nan_to_num(kpts)

        oob_mask = out_of_bounds_keypoints(kpts, self.img_size)
        if np.sum(oob_mask) > 0:
            kpts[oob_mask] = 0
        kpts = kpts.astype(int)

        zero_matrix = np.zeros((batch_size, size[0], size[1], self.num_channels))

        def _get_condition_matrix(zero_matrix, kpts):
            for i, pose in enumerate(kpts):
                x, y, vis = pose.T
                mask = vis > 0
                x_masked, y_masked, colors_masked = x[mask], y[mask], self.colors[mask]
                zero_matrix[i, y_masked - 1, x_masked - 1] = colors_masked
            return zero_matrix

        def _get_condition_matrix_optim(zero_matrix, kpts):
            x, y = np.array(kpts).T
            mask = (0 < x) & (x < zero_matrix.shape[2]) & (0 < y) & (y < zero_matrix.shape[1])
            colors_masked = np.repeat(self.colors[:, None, :], len(zero_matrix), 1) * np.repeat(mask[:, :, None], 3, 2)
            kpt_indices = np.stack([x.T, y.T]).transpose(1, 2, 0)
            batch_indices = np.repeat(np.arange(len(zero_matrix))[:, None, None], self.num_joints, axis=1)
            kpt_input = np.concatenate([batch_indices, kpt_indices], dtype=int, axis=2)
            zero_matrix[kpt_input[..., 0], kpt_input[..., 2] - 1, kpt_input[..., 1] - 1] = colors_masked.transpose(
                1, 0, 2
            )
            return zero_matrix

        condition = _get_condition_matrix(zero_matrix, kpts)
        # condition = _get_condition_matrix_optim(zero_matrix, kpts)

        for i in range(batch_size):
            condition_heatmap = self.blur_heatmap(condition[i])
            condition[i] = condition_heatmap
        # condition = self.blur_heatmap_batch(torch.from_numpy(condition))

        return condition

    def get_colors_from_cmap(self, cmap_name, num_colors):
        cmap = plt.get_cmap(cmap_name)
        colors_float = [cmap(i) for i in np.linspace(0, 256, num_colors, dtype=int)]
        colors = [(int(r * 255), int(g * 255), int(b * 255)) for r, g, b, _ in colors_float]
        return colors

__call__ 

__call__(keypoints: ndarray, size: tuple[int, int]) -> np.ndarray

Parameters:

Name Type Description Default

keypoints

ndarray

batch of keypoints to encode with shape (batch_size, num_joints, 2)

 required

size

tuple[int, int]

the (height, width) of the heatmap in which the keypoints should be encoded

 required

Returns:

Type Description
ndarray

encoded keypoints with shape (batch_size, num_joints, height, width, 3)
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
175
176
177
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
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
def __call__(self, keypoints: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """
    Args:
        keypoints: batch of keypoints to encode with shape (batch_size, num_joints, 2)
        size: the (height, width) of the heatmap in which the keypoints should be encoded

    Returns:
        encoded keypoints with shape (batch_size, num_joints, height, width, 3)
    """

    batch_size, num_kpts, _ = keypoints.shape

    if not num_kpts == len(self.colors):
        raise ValueError(
            f"Cannot encode the keypoints. Initialized with {len(self.colors)} "
            f"colors, but there are {num_kpts} to encode"
        )

    # kpts = keypoints.detach().numpy()
    kpts = keypoints.copy()
    kpts[keypoints[..., 2] <= 0] = 0

    # Mark keypoints as visible, remove NaNs
    kpts[kpts[..., 2] > 0, 2] = 2
    kpts = np.nan_to_num(kpts)

    oob_mask = out_of_bounds_keypoints(kpts, self.img_size)
    if np.sum(oob_mask) > 0:
        kpts[oob_mask] = 0
    kpts = kpts.astype(int)

    zero_matrix = np.zeros((batch_size, size[0], size[1], self.num_channels))

    def _get_condition_matrix(zero_matrix, kpts):
        for i, pose in enumerate(kpts):
            x, y, vis = pose.T
            mask = vis > 0
            x_masked, y_masked, colors_masked = x[mask], y[mask], self.colors[mask]
            zero_matrix[i, y_masked - 1, x_masked - 1] = colors_masked
        return zero_matrix

    def _get_condition_matrix_optim(zero_matrix, kpts):
        x, y = np.array(kpts).T
        mask = (0 < x) & (x < zero_matrix.shape[2]) & (0 < y) & (y < zero_matrix.shape[1])
        colors_masked = np.repeat(self.colors[:, None, :], len(zero_matrix), 1) * np.repeat(mask[:, :, None], 3, 2)
        kpt_indices = np.stack([x.T, y.T]).transpose(1, 2, 0)
        batch_indices = np.repeat(np.arange(len(zero_matrix))[:, None, None], self.num_joints, axis=1)
        kpt_input = np.concatenate([batch_indices, kpt_indices], dtype=int, axis=2)
        zero_matrix[kpt_input[..., 0], kpt_input[..., 2] - 1, kpt_input[..., 1] - 1] = colors_masked.transpose(
            1, 0, 2
        )
        return zero_matrix

    condition = _get_condition_matrix(zero_matrix, kpts)
    # condition = _get_condition_matrix_optim(zero_matrix, kpts)

    for i in range(batch_size):
        condition_heatmap = self.blur_heatmap(condition[i])
        condition[i] = condition_heatmap
    # condition = self.blur_heatmap_batch(torch.from_numpy(condition))

    return condition

__init__ 

__init__(colors: list[tuple[int, int, int]] | None = None, **kwargs) -> None

Parameters:

Name Type Description Default

colors

list[tuple[int, int, int]] | None

the color to use for each keypoint

 None
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
161
162
163
164
165
166
167
168
169
def __init__(self, colors: list[tuple[int, int, int]] | None = None, **kwargs) -> None:
    """
    Args:
        colors: the color to use for each keypoint
    """
    super().__init__(**kwargs)
    if colors is None:
        colors = self.get_colors_from_cmap("rainbow", self.num_joints)
    self.colors = np.array(colors)

StackedKeypointEncoder

Bases: BaseKeypointEncoder

Encodes keypoints into heatmaps, where each.

Modified from BUCTD/data/JointsDataset, get_stacked_condition

Methods:

Name Description
__call__

Args:
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
 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
@KEYPOINT_ENCODERS.register_module
class StackedKeypointEncoder(BaseKeypointEncoder):
    """Encodes keypoints into heatmaps, where each.

    Modified from BUCTD/data/JointsDataset, get_stacked_condition
    """

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)

    @property
    def num_channels(self):
        return self.num_joints

    def __call__(self, keypoints: np.ndarray, size: tuple[int, int]) -> np.ndarray:
        """
        Args:
            keypoints: the keypoints to encode
            size: the (height, width) of the heatmap in which the keypoints should
                be encoded

        Returns:
            the encoded keypoints
        """

        batch_size, _, _ = keypoints.shape

        kpts = keypoints.copy()
        kpts[keypoints[..., 2] <= 0] = 0

        # Mark keypoints as visible, remove NaNs
        kpts[kpts[..., 2] > 0, 2] = 2
        kpts = np.nan_to_num(kpts)

        oob_mask = out_of_bounds_keypoints(kpts, self.img_size)
        if np.sum(oob_mask) > 0:
            kpts[oob_mask] = 0
        kpts = kpts.astype(int)

        zero_matrix = np.zeros((batch_size, size[0], size[1], self.num_channels))

        def _get_condition_matrix(zero_matrix, kpts):
            for i, pose in enumerate(kpts):
                x, y, vis = pose.T
                mask = vis > 0
                x_masked, y_masked, joint_inds_masked = (
                    x[mask],
                    y[mask],
                    np.arange(self.num_joints)[mask],
                )
                zero_matrix[i, y_masked - 1, x_masked - 1, joint_inds_masked] = 255
            return zero_matrix

        condition = _get_condition_matrix(zero_matrix, kpts)

        for i in range(batch_size):
            condition_heatmap = self.blur_heatmap(condition[i])
            condition[i] = condition_heatmap

        return condition

__call__ 

__call__(keypoints: ndarray, size: tuple[int, int]) -> np.ndarray

Parameters:

Name Type Description Default

keypoints

ndarray

the keypoints to encode

 required

size

tuple[int, int]

the (height, width) of the heatmap in which the keypoints should be encoded

 required

Returns:

Type Description
ndarray

the encoded keypoints
Source code in deeplabcut/pose_estimation_pytorch/models/modules/kpt_encoders.py 
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
def __call__(self, keypoints: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    """
    Args:
        keypoints: the keypoints to encode
        size: the (height, width) of the heatmap in which the keypoints should
            be encoded

    Returns:
        the encoded keypoints
    """

    batch_size, _, _ = keypoints.shape

    kpts = keypoints.copy()
    kpts[keypoints[..., 2] <= 0] = 0

    # Mark keypoints as visible, remove NaNs
    kpts[kpts[..., 2] > 0, 2] = 2
    kpts = np.nan_to_num(kpts)

    oob_mask = out_of_bounds_keypoints(kpts, self.img_size)
    if np.sum(oob_mask) > 0:
        kpts[oob_mask] = 0
    kpts = kpts.astype(int)

    zero_matrix = np.zeros((batch_size, size[0], size[1], self.num_channels))

    def _get_condition_matrix(zero_matrix, kpts):
        for i, pose in enumerate(kpts):
            x, y, vis = pose.T
            mask = vis > 0
            x_masked, y_masked, joint_inds_masked = (
                x[mask],
                y[mask],
                np.arange(self.num_joints)[mask],
            )
            zero_matrix[i, y_masked - 1, x_masked - 1, joint_inds_masked] = 255
        return zero_matrix

    condition = _get_condition_matrix(zero_matrix, kpts)

    for i in range(batch_size):
        condition_heatmap = self.blur_heatmap(condition[i])
        condition[i] = condition_heatmap

    return condition