|
from typing import Dict, List |
|
|
|
import numpy as np |
|
import torch |
|
from torch import Tensor |
|
|
|
import mGPT.utils.geometry_conver as geometry_conver |
|
|
|
|
|
def lengths_to_mask(lengths: List[int], |
|
device: torch.device, |
|
max_len: int = None) -> Tensor: |
|
lengths = torch.tensor(lengths, device=device) |
|
max_len = max_len if max_len else max(lengths) |
|
mask = torch.arange(max_len, device=device).expand( |
|
len(lengths), max_len) < lengths.unsqueeze(1) |
|
return mask |
|
|
|
|
|
def detach_to_numpy(tensor): |
|
return tensor.detach().cpu().numpy() |
|
|
|
|
|
def remove_padding(tensors, lengths): |
|
return [ |
|
tensor[:tensor_length] |
|
for tensor, tensor_length in zip(tensors, lengths) |
|
] |
|
|
|
|
|
def nfeats_of(rottype): |
|
if rottype in ["rotvec", "axisangle"]: |
|
return 3 |
|
elif rottype in ["rotquat", "quaternion"]: |
|
return 4 |
|
elif rottype in ["rot6d", "6drot", "rotation6d"]: |
|
return 6 |
|
elif rottype in ["rotmat"]: |
|
return 9 |
|
else: |
|
return TypeError("This rotation type doesn't have features.") |
|
|
|
|
|
def axis_angle_to(newtype, rotations): |
|
if newtype in ["matrix"]: |
|
rotations = geometry_conver.axis_angle_to_matrix(rotations) |
|
return rotations |
|
elif newtype in ["rotmat"]: |
|
rotations = geometry_conver.axis_angle_to_matrix(rotations) |
|
rotations = matrix_to("rotmat", rotations) |
|
return rotations |
|
elif newtype in ["rot6d", "6drot", "rotation6d"]: |
|
rotations = geometry_conver.axis_angle_to_matrix(rotations) |
|
rotations = matrix_to("rot6d", rotations) |
|
return rotations |
|
elif newtype in ["rotquat", "quaternion"]: |
|
rotations = geometry_conver.axis_angle_to_quaternion(rotations) |
|
return rotations |
|
elif newtype in ["rotvec", "axisangle"]: |
|
return rotations |
|
else: |
|
raise NotImplementedError |
|
|
|
|
|
def matrix_to(newtype, rotations): |
|
if newtype in ["matrix"]: |
|
return rotations |
|
if newtype in ["rotmat"]: |
|
rotations = rotations.reshape((*rotations.shape[:-2], 9)) |
|
return rotations |
|
elif newtype in ["rot6d", "6drot", "rotation6d"]: |
|
rotations = geometry_conver.matrix_to_rotation_6d(rotations) |
|
return rotations |
|
elif newtype in ["rotquat", "quaternion"]: |
|
rotations = geometry_conver.matrix_to_quaternion(rotations) |
|
return rotations |
|
elif newtype in ["rotvec", "axisangle"]: |
|
rotations = geometry_conver.matrix_to_axis_angle(rotations) |
|
return rotations |
|
else: |
|
raise NotImplementedError |
|
|
|
|
|
def to_matrix(oldtype, rotations): |
|
if oldtype in ["matrix"]: |
|
return rotations |
|
if oldtype in ["rotmat"]: |
|
rotations = rotations.reshape((*rotations.shape[:-2], 3, 3)) |
|
return rotations |
|
elif oldtype in ["rot6d", "6drot", "rotation6d"]: |
|
rotations = geometry_conver.rotation_6d_to_matrix(rotations) |
|
return rotations |
|
elif oldtype in ["rotquat", "quaternion"]: |
|
rotations = geometry_conver.quaternion_to_matrix(rotations) |
|
return rotations |
|
elif oldtype in ["rotvec", "axisangle"]: |
|
rotations = geometry_conver.axis_angle_to_matrix(rotations) |
|
return rotations |
|
else: |
|
raise NotImplementedError |
|
|
|
|
|
|
|
def subsample(num_frames, last_framerate, new_framerate): |
|
step = int(last_framerate / new_framerate) |
|
assert step >= 1 |
|
frames = np.arange(0, num_frames, step) |
|
return frames |
|
|
|
|
|
|
|
def upsample(motion, last_framerate, new_framerate): |
|
step = int(new_framerate / last_framerate) |
|
assert step >= 1 |
|
|
|
|
|
alpha = np.linspace(0, 1, step + 1) |
|
last = np.einsum("l,...->l...", 1 - alpha, motion[:-1]) |
|
new = np.einsum("l,...->l...", alpha, motion[1:]) |
|
|
|
chuncks = (last + new)[:-1] |
|
output = np.concatenate(chuncks.swapaxes(1, 0)) |
|
|
|
output = np.concatenate((output, motion[[-1]])) |
|
return output |
|
|
|
|
|
if __name__ == "__main__": |
|
motion = np.arange(105) |
|
submotion = motion[subsample(len(motion), 100.0, 12.5)] |
|
newmotion = upsample(submotion, 12.5, 100) |
|
|
|
print(newmotion) |
|
|
