Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,642 Bytes
352b049 |
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 |
import numpy as np
import torch
from .base import LANDMARKS
def convert2humanml(
motion: torch.tensor,
dataset_landmarks: list[str],
joint_to_humanml_names: list[tuple[str, str]],
) -> np.ndarray:
"""
Convert skeletons to human ml.
Args:
skeletons: Skeletons.
dataset_landmarks: Landmarks of the dataset.
joint_to_humanml_names: Conversion names.
Returns:
new_motion: Skeletons in human ml.
"""
new_motion = torch.zeros((*motion.shape[:-2], len(LANDMARKS), 3), device=motion.device)
indices = torch.tensor(
[
(LANDMARKS.index(hml_name), dataset_landmarks.index(nba_name))
for hml_name, nba_name in joint_to_humanml_names
]
)
new_motion[..., indices[:, 0], :] = motion[..., indices[:, 1], :]
# Infer missing landmarks
covered_landmarks = [hml_landmark for hml_landmark, _ in joint_to_humanml_names]
root_index, left_hip_index, right_hip_index = (
LANDMARKS.index("pelvis"),
LANDMARKS.index("left_hip"),
LANDMARKS.index("right_hip"),
)
if "pelvis" not in covered_landmarks:
assert (
"left_hip" in covered_landmarks and "right_hip" in covered_landmarks
), "pelvis could not be inferred from left_hip and right_hip"
new_motion[..., root_index, :] = (
new_motion[..., left_hip_index, :] + new_motion[..., right_hip_index, :]
) / 2
neck_index, left_shoulder_index, right_shoulder_index = (
LANDMARKS.index("neck"),
LANDMARKS.index("left_shoulder"),
LANDMARKS.index("right_shoulder"),
)
if "neck" not in covered_landmarks:
assert (
"left_shoulder" in covered_landmarks
and "right_shoulder" in covered_landmarks
), "neck could not be inferred from left_shoulder and right_shoulder"
new_motion[..., neck_index, :] = (
new_motion[..., left_shoulder_index, :]
+ new_motion[..., right_shoulder_index, :]
) / 2
spine_1_index, spine_2_index, spine_3_index = (
LANDMARKS.index("spine1"),
LANDMARKS.index("spine2"),
LANDMARKS.index("spine3"),
)
if "spine1" not in covered_landmarks:
new_motion[..., spine_1_index, :] = (
0.75 * new_motion[..., root_index, :]
+ 0.25 * new_motion[..., neck_index, :]
)
if "spine2" not in covered_landmarks:
new_motion[..., spine_2_index, :] = (
0.5 * new_motion[..., root_index, :] + 0.5 * new_motion[..., neck_index, :]
)
if "spine3" not in covered_landmarks:
new_motion[..., spine_3_index, :] = (
0.25 * new_motion[..., root_index, :]
+ 0.75 * new_motion[..., neck_index, :]
)
left_collar_index, right_collar_index = LANDMARKS.index(
"left_collar"
), LANDMARKS.index("right_collar")
if "left_collar" not in covered_landmarks:
assert (
"left_shoulder" in covered_landmarks
), "left_collar could not be inferred from left_shoulder"
new_motion[..., left_collar_index, :] = (
0.25 * new_motion[..., left_shoulder_index, :]
+ 0.75 * new_motion[..., neck_index, :]
)
if "right_collar" not in covered_landmarks:
assert (
"right_shoulder" in covered_landmarks
), "right_collar could not be inferred from right_shoulder"
new_motion[..., right_collar_index, :] = (
0.25 * new_motion[..., right_shoulder_index, :]
+ 0.75 * new_motion[..., neck_index, :]
)
return new_motion
|