/
image_pose_lifting.py
280 lines (237 loc) · 12.1 KB
/
image_pose_lifting.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
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
243
244
245
246
247
248
249
250
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
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Tuple, Union
import numpy as np
from mmpose.registry import KEYPOINT_CODECS
from .base import BaseKeypointCodec
@KEYPOINT_CODECS.register_module()
class ImagePoseLifting(BaseKeypointCodec):
r"""Generate keypoint coordinates for pose lifter.
Note:
- instance number: N
- keypoint number: K
- keypoint dimension: D
- pose-lifitng target dimension: C
Args:
num_keypoints (int): The number of keypoints in the dataset.
root_index (Union[int, List]): Root keypoint index in the pose.
remove_root (bool): If true, remove the root keypoint from the pose.
Default: ``False``.
save_index (bool): If true, store the root position separated from the
original pose. Default: ``False``.
reshape_keypoints (bool): If true, reshape the keypoints into shape
(-1, N). Default: ``True``.
concat_vis (bool): If true, concat the visibility item of keypoints.
Default: ``False``.
keypoints_mean (np.ndarray, optional): Mean values of keypoints
coordinates in shape (K, D).
keypoints_std (np.ndarray, optional): Std values of keypoints
coordinates in shape (K, D).
target_mean (np.ndarray, optional): Mean values of pose-lifitng target
coordinates in shape (K, C).
target_std (np.ndarray, optional): Std values of pose-lifitng target
coordinates in shape (K, C).
"""
auxiliary_encode_keys = {'lifting_target', 'lifting_target_visible'}
instance_mapping_table = dict(
lifting_target='lifting_target',
lifting_target_visible='lifting_target_visible',
)
label_mapping_table = dict(
trajectory_weights='trajectory_weights',
lifting_target_label='lifting_target_label',
lifting_target_weight='lifting_target_weight')
def __init__(self,
num_keypoints: int,
root_index: Union[int, List] = 0,
remove_root: bool = False,
save_index: bool = False,
reshape_keypoints: bool = True,
concat_vis: bool = False,
keypoints_mean: Optional[np.ndarray] = None,
keypoints_std: Optional[np.ndarray] = None,
target_mean: Optional[np.ndarray] = None,
target_std: Optional[np.ndarray] = None,
additional_encode_keys: Optional[List[str]] = None):
super().__init__()
self.num_keypoints = num_keypoints
if isinstance(root_index, int):
root_index = [root_index]
self.root_index = root_index
self.remove_root = remove_root
self.save_index = save_index
self.reshape_keypoints = reshape_keypoints
self.concat_vis = concat_vis
if keypoints_mean is not None:
assert keypoints_std is not None, 'keypoints_std is None'
keypoints_mean = np.array(
keypoints_mean,
dtype=np.float32).reshape(1, num_keypoints, -1)
keypoints_std = np.array(
keypoints_std, dtype=np.float32).reshape(1, num_keypoints, -1)
assert keypoints_mean.shape == keypoints_std.shape, (
f'keypoints_mean.shape {keypoints_mean.shape} != '
f'keypoints_std.shape {keypoints_std.shape}')
if target_mean is not None:
assert target_std is not None, 'target_std is None'
target_dim = num_keypoints - 1 if remove_root else num_keypoints
target_mean = np.array(
target_mean, dtype=np.float32).reshape(1, target_dim, -1)
target_std = np.array(
target_std, dtype=np.float32).reshape(1, target_dim, -1)
assert target_mean.shape == target_std.shape, (
f'target_mean.shape {target_mean.shape} != '
f'target_std.shape {target_std.shape}')
self.keypoints_mean = keypoints_mean
self.keypoints_std = keypoints_std
self.target_mean = target_mean
self.target_std = target_std
if additional_encode_keys is not None:
self.auxiliary_encode_keys.update(additional_encode_keys)
def encode(self,
keypoints: np.ndarray,
keypoints_visible: Optional[np.ndarray] = None,
lifting_target: Optional[np.ndarray] = None,
lifting_target_visible: Optional[np.ndarray] = None) -> dict:
"""Encoding keypoints from input image space to normalized space.
Args:
keypoints (np.ndarray): Keypoint coordinates in shape (N, K, D).
keypoints_visible (np.ndarray, optional): Keypoint visibilities in
shape (N, K).
lifting_target (np.ndarray, optional): 3d target coordinate in
shape (T, K, C).
lifting_target_visible (np.ndarray, optional): Target coordinate in
shape (T, K, ).
Returns:
encoded (dict): Contains the following items:
- keypoint_labels (np.ndarray): The processed keypoints in
shape like (N, K, D) or (K * D, N).
- keypoint_labels_visible (np.ndarray): The processed
keypoints' weights in shape (N, K, ) or (N-1, K, ).
- lifting_target_label: The processed target coordinate in
shape (K, C) or (K-1, C).
- lifting_target_weight (np.ndarray): The target weights in
shape (K, ) or (K-1, ).
- trajectory_weights (np.ndarray): The trajectory weights in
shape (K, ).
- target_root (np.ndarray): The root coordinate of target in
shape (C, ).
In addition, there are some optional items it may contain:
- target_root (np.ndarray): The root coordinate of target in
shape (C, ). Exists if ``zero_center`` is ``True``.
- target_root_removed (bool): Indicate whether the root of
pose-lifitng target is removed. Exists if
``remove_root`` is ``True``.
- target_root_index (int): An integer indicating the index of
root. Exists if ``remove_root`` and ``save_index``
are ``True``.
"""
if keypoints_visible is None:
keypoints_visible = np.ones(keypoints.shape[:2], dtype=np.float32)
if lifting_target is None:
lifting_target = [keypoints[0]]
# set initial value for `lifting_target_weight`
# and `trajectory_weights`
if lifting_target_visible is None:
lifting_target_visible = np.ones(
lifting_target.shape[:-1], dtype=np.float32)
lifting_target_weight = lifting_target_visible
trajectory_weights = (1 / lifting_target[:, 2])
else:
valid = lifting_target_visible > 0.5
lifting_target_weight = np.where(valid, 1., 0.).astype(np.float32)
trajectory_weights = lifting_target_weight
encoded = dict()
# Zero-center the target pose around a given root keypoint
assert (lifting_target.ndim >= 2 and
lifting_target.shape[-2] > max(self.root_index)), \
f'Got invalid joint shape {lifting_target.shape}'
root = np.mean(
lifting_target[..., self.root_index, :], axis=-2, dtype=np.float32)
lifting_target_label = lifting_target - root[np.newaxis, ...]
if self.remove_root and len(self.root_index) == 1:
root_index = self.root_index[0]
lifting_target_label = np.delete(
lifting_target_label, root_index, axis=-2)
lifting_target_visible = np.delete(
lifting_target_visible, root_index, axis=-2)
assert lifting_target_weight.ndim in {
2, 3
}, (f'lifting_target_weight.ndim {lifting_target_weight.ndim} '
'is not in {2, 3}')
axis_to_remove = -2 if lifting_target_weight.ndim == 3 else -1
lifting_target_weight = np.delete(
lifting_target_weight, root_index, axis=axis_to_remove)
# Add a flag to avoid latter transforms that rely on the root
# joint or the original joint index
encoded['target_root_removed'] = True
# Save the root index which is necessary to restore the global pose
if self.save_index:
encoded['target_root_index'] = root_index
# Normalize the 2D keypoint coordinate with mean and std
keypoint_labels = keypoints.copy()
if self.keypoints_mean is not None:
assert self.keypoints_mean.shape[1:] == keypoints.shape[1:], (
f'self.keypoints_mean.shape[1:] {self.keypoints_mean.shape[1:]} ' # noqa
f'!= keypoints.shape[1:] {keypoints.shape[1:]}')
encoded['keypoints_mean'] = self.keypoints_mean.copy()
encoded['keypoints_std'] = self.keypoints_std.copy()
keypoint_labels = (keypoint_labels -
self.keypoints_mean) / self.keypoints_std
if self.target_mean is not None:
assert self.target_mean.shape == lifting_target_label.shape, (
f'self.target_mean.shape {self.target_mean.shape} '
f'!= lifting_target_label.shape {lifting_target_label.shape}' # noqa
)
encoded['target_mean'] = self.target_mean.copy()
encoded['target_std'] = self.target_std.copy()
lifting_target_label = (lifting_target_label -
self.target_mean) / self.target_std
# Generate reshaped keypoint coordinates
assert keypoint_labels.ndim in {
2, 3
}, (f'keypoint_labels.ndim {keypoint_labels.ndim} is not in {2, 3}')
if keypoint_labels.ndim == 2:
keypoint_labels = keypoint_labels[None, ...]
if self.concat_vis:
keypoints_visible_ = keypoints_visible
if keypoints_visible.ndim == 2:
keypoints_visible_ = keypoints_visible[..., None]
keypoint_labels = np.concatenate(
(keypoint_labels, keypoints_visible_), axis=2)
if self.reshape_keypoints:
N = keypoint_labels.shape[0]
keypoint_labels = keypoint_labels.transpose(1, 2, 0).reshape(-1, N)
encoded['keypoint_labels'] = keypoint_labels
encoded['keypoint_labels_visible'] = keypoints_visible
encoded['lifting_target_label'] = lifting_target_label
encoded['lifting_target_weight'] = lifting_target_weight
encoded['trajectory_weights'] = trajectory_weights
encoded['target_root'] = root
return encoded
def decode(self,
encoded: np.ndarray,
target_root: Optional[np.ndarray] = None
) -> Tuple[np.ndarray, np.ndarray]:
"""Decode keypoint coordinates from normalized space to input image
space.
Args:
encoded (np.ndarray): Coordinates in shape (N, K, C).
target_root (np.ndarray, optional): The target root coordinate.
Default: ``None``.
Returns:
keypoints (np.ndarray): Decoded coordinates in shape (N, K, C).
scores (np.ndarray): The keypoint scores in shape (N, K).
"""
keypoints = encoded.copy()
if self.target_mean is not None and self.target_std is not None:
assert self.target_mean.shape == keypoints.shape, (
f'self.target_mean.shape {self.target_mean.shape} '
f'!= keypoints.shape {keypoints.shape}')
keypoints = keypoints * self.target_std + self.target_mean
if target_root is not None and target_root.size > 0:
keypoints = keypoints + target_root
if self.remove_root and len(self.root_index) == 1:
keypoints = np.insert(
keypoints, self.root_index, target_root, axis=1)
scores = np.ones(keypoints.shape[:-1], dtype=np.float32)
return keypoints, scores