Spaces:
Runtime error
Runtime error
File size: 4,271 Bytes
2d5fdd1 |
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 |
import numpy as np
# import librosa
from pathlib import Path
import json
import os.path
import sys
import argparse
import pickle
import torch
THIS_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.abspath(os.path.join(os.path.join(THIS_DIR, os.pardir), os.pardir))
DATA_DIR = os.path.join(ROOT_DIR, 'data')
EXTRACT_DIR = os.path.join(DATA_DIR, 'extracted_data')
if not os.path.isdir(DATA_DIR):
os.mkdir(DATA_DIR)
if not os.path.isdir(EXTRACT_DIR):
os.mkdir(EXTRACT_DIR)
sys.path.append(ROOT_DIR)
from utils import distribute_tasks
parser = argparse.ArgumentParser(description="Preprocess motion data")
parser.add_argument("data_path", type=str, help="Directory contining Beat Saber level folders")
parser.add_argument("--replace_existing", action="store_true")
args = parser.parse_args()
# makes arugments into global variables of the same name, used later in the code
globals().update(vars(args))
data_path = Path(data_path)
from scipy.spatial.transform import Rotation as R
from typing import NewType, Union, Optional
Tensor = NewType('Tensor', torch.Tensor)
def rot_mat_to_euler(rot_mats):
# Calculates rotation matrix to euler angles
# Careful for extreme cases of eular angles like [0.0, pi, 0.0]
sy = torch.sqrt(rot_mats[:, 0, 0] * rot_mats[:, 0, 0] +
rot_mats[:, 1, 0] * rot_mats[:, 1, 0])
return torch.atan2(-rot_mats[:, 2, 0], sy)
def batch_rodrigues(
rot_vecs: Tensor,
epsilon: float = 1e-8,
) -> Tensor:
''' Calculates the rotation matrices for a batch of rotation vectors
Parameters
----------
rot_vecs: torch.tensor Nx3
array of N axis-angle vectors
Returns
-------
R: torch.tensor Nx3x3
The rotation matrices for the given axis-angle parameters
'''
batch_size = rot_vecs.shape[0]
device, dtype = rot_vecs.device, rot_vecs.dtype
angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
rot_dir = rot_vecs / angle
cos = torch.unsqueeze(torch.cos(angle), dim=1)
sin = torch.unsqueeze(torch.sin(angle), dim=1)
# Bx1 arrays
rx, ry, rz = torch.split(rot_dir, 1, dim=1)
K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)
zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
.view((batch_size, 3, 3))
ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
return rot_mat
def get_rot_matrices_from_euler(joint_traj):
return np.stack([np.concatenate([R.from_euler('xyz',euler_angles).as_matrix().flatten() for euler_angles in np.array(joint_angles).reshape(-1,3)]) for joint_angles in joint_traj])
def get_rot_matrices_from_axis_angle(joint_traj):
# return np.stack([np.concatenate([R.from_euler('xyz',euler_angles).as_matrix().flatten() for euler_angles in np.array(joint_angles).reshape(-1,3)]) for joint_angles in joint_traj])
L = len(joint_traj)
ident = torch.eye(3, dtype=torch.float32)
rot_mats = batch_rodrigues(torch.from_numpy(joint_traj).view(-1, 3)).view(
[L,-1, 3, 3])
pose_feature = (rot_mats - ident).view([L, -1])
return pose_feature.numpy()
def get_features(motion_data):
joint_angle_feats = get_rot_matrices_from_axis_angle((motion_data['smpl_poses']))
return np.concatenate([joint_angle_feats,motion_data['smpl_trans']],1)
## distributing tasks accross nodes ##
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
print(rank)
candidate_motion_files = sorted(data_path.glob('**/*.pkl'), key=lambda path: path.parent.__str__())
tasks = distribute_tasks(candidate_motion_files,rank,size)
for i in tasks:
path = candidate_motion_files[i]
motion_file_path = path.__str__()
features_file = motion_file_path+"_"+"joint_angles_mats"+".npy"
if replace_existing or not os.path.isfile(features_file):
motion_data = pickle.load(open(path,"rb"))
features = get_features(motion_data)
print(features.shape)
features = ResampleLinear1D(features,features.shape[0]*2)
print(features.shape)
np.save(features_file,features)
|